Source File
	serialize.go

Belonging Package
	github.com/freeeve/gochickpeas

// RCPG serialization for Snapshot: conversion to and from the rcpg
// package's on-disk data model. Lazily built indexes are derived data and
// are not serialized; they rebuild on first use after reading.

package chickpeas

import (
	"bufio"
	"io"
	"os"
	"sort"

	"github.com/RoaringBitmap/roaring/v2"
	"github.com/freeeve/gochickpeas/internal/bitset"
	"github.com/freeeve/gochickpeas/nodeset"
	"github.com/freeeve/gochickpeas/rcpg"
)

// computeInToOutFromCSR recomputes the incoming -> outgoing CSR position
// map from the adjacency arrays alone (used when a snapshot is
// reconstructed on deserialization). It pairs the k-th u->v rel of a given
// type in the outgoing CSR with the k-th such rel in the incoming CSR; both
// CSRs preserve insertion order within a group, so the pairing matches the
// original rels.
//
// Implementation: stable LSD counting sorts bring both sides' position
// arrays into the identical (src, dst, type) group order -- stability
// keeps each side's within-group insertion order, so the i-th entries of
// the two sorted arrays are the same relationship and zip directly. This
// replaces the former per-group map (one entry plus a slice per distinct
// rel key -- the dominant load-time cost at tens of millions of rels)
// with a handful of flat arrays.
func computeInToOutFromCSR(outOffsets []uint32, outNbrs []NodeID, outTypes []RelType,
	inOffsets []uint32, inNbrs []NodeID, inTypes []RelType) []uint32 {
	m := len(outNbrs)
	inToOut := make([]uint32, m)
	if m == 0 {
		return inToOut
	}
	n := max(len(inOffsets)-1, 0)
	outSrc := ownersOf(outOffsets, m)
	inDst := ownersOf(inOffsets, m)
	maxType := RelType(0)
	for _, t := range outTypes {
		maxType = max(maxType, t)
	}

	outPos := ascending(m)
	inPos := ascending(m)
	tmp := make([]uint32, m)
	counts := make([]uint32, max(n, int(maxType)+1)+1)
	// LSD: least-significant key first, so the LAST pass's key is most
	// significant; both sides end in the same (src, dst, type) order.
	countingPass(outPos, tmp, counts, int(maxType)+1, func(p uint32) uint32 { return uint32(outTypes[p]) })
	countingPass(outPos, tmp, counts, n, func(p uint32) uint32 { return uint32(outNbrs[p]) })
	countingPass(outPos, tmp, counts, n, func(p uint32) uint32 { return outSrc[p] })
	countingPass(inPos, tmp, counts, int(maxType)+1, func(p uint32) uint32 { return uint32(inTypes[p]) })
	countingPass(inPos, tmp, counts, n, func(p uint32) uint32 { return inDst[p] })
	countingPass(inPos, tmp, counts, n, func(p uint32) uint32 { return uint32(inNbrs[p]) })

	for i := range m {
		inToOut[inPos[i]] = outPos[i]
	}
	return inToOut
}

// ownersOf expands a CSR offsets array to a per-position owner-node array.
func ownersOf(offsets []uint32, m int) []uint32 {
	owners := make([]uint32, m)
	for u := 0; u+1 < len(offsets); u++ {
		for p := offsets[u]; p < offsets[u+1]; p++ {
			owners[p] = uint32(u)
		}
	}
	return owners
}

// ascending is the identity position permutation.
func ascending(m int) []uint32 {
	out := make([]uint32, m)
	for i := range m {
		out[i] = uint32(i)
	}
	return out
}

// countingPass stably reorders positions in place by key(p) in [0, nKeys),
// staging through tmp.
func countingPass(positions, tmp, counts []uint32, nKeys int, key func(uint32) uint32) {
	clear(counts[:nKeys+1])
	for _, p := range positions {
		counts[key(p)+1]++
	}
	for k := 1; k <= nKeys; k++ {
		counts[k] += counts[k-1]
	}
	for _, p := range positions {
		k := key(p)
		tmp[counts[k]] = p
		counts[k]++
	}
	copy(positions, tmp)
}

func columnToData(c Column) rcpg.ColumnData {
	switch col := c.(type) {
	case denseI64Col:
		return rcpg.DenseI64(col)
	case denseI64NarrowCol:
		// Serialization stays width-agnostic: materialize the logical
		// values; the reader re-narrows on load.
		out := make(rcpg.DenseI64, col.Len())
		for i := range out {
			out[i] = col.at(uint32(i))
		}
		return out
	case denseI64BitCol:
		out := make(rcpg.DenseI64, col.Len())
		for i := range out {
			out[i] = col.at(uint32(i))
		}
		return out
	case denseF64Col:
		return rcpg.DenseF64(col)
	case denseBoolCol:
		return rcpg.DenseBool{Bits: col.bits.ToPackedLSB(), Len: uint32(col.bits.Len())}
	case denseStrCol:
		return rcpg.DenseStr(col)
	case sparseI64Col:
		out := make(rcpg.SparseI64, len(col.ids))
		for i, id := range col.ids {
			out[i] = rcpg.I64Entry{ID: id, Val: col.vals[i]}
		}
		return out
	case sparseI64NarrowCol:
		out := make(rcpg.SparseI64, len(col.ids))
		for i, id := range col.ids {
			out[i] = rcpg.I64Entry{ID: id, Val: col.at(i)}
		}
		return out
	case rankI64NarrowCol:
		// Like rankI64Col: rank/select is in-memory only, serialized as
		// the equivalent sparse column.
		out := make(rcpg.SparseI64, 0, col.n())
		for pos, v := range col.Entries() {
			x, _ := v.I64()
			out = append(out, rcpg.I64Entry{ID: pos, Val: x})
		}
		return out
	case sparseF64Col:
		out := make(rcpg.SparseF64, len(col.ids))
		for i, id := range col.ids {
			out[i] = rcpg.F64Entry{ID: id, Val: col.vals[i]}
		}
		return out
	case sparseBoolCol:
		out := make(rcpg.SparseBool, len(col.ids))
		for i, id := range col.ids {
			out[i] = rcpg.BoolEntry{ID: id, Val: col.vals[i]}
		}
		return out
	case sparseStrCol:
		out := make(rcpg.SparseStr, len(col.ids))
		for i, id := range col.ids {
			out[i] = rcpg.StrEntry{ID: id, Val: col.vals[i]}
		}
		return out
	// Rank/select is an in-memory layout only; it serializes as the
	// equivalent sparse column (matching the Rust codec).
	case rankI64Col:
		out := make(rcpg.SparseI64, 0, len(col.vals))
		for pos, v := range col.Entries() {
			x, _ := v.I64()
			out = append(out, rcpg.I64Entry{ID: pos, Val: x})
		}
		return out
	case rankF64Col:
		out := make(rcpg.SparseF64, 0, len(col.vals))
		for pos, v := range col.Entries() {
			x, _ := v.F64()
			out = append(out, rcpg.F64Entry{ID: pos, Val: x})
		}
		return out
	case rankBoolCol:
		out := make(rcpg.SparseBool, 0, col.vals.Len())
		for pos, v := range col.Entries() {
			x, _ := v.Bool()
			out = append(out, rcpg.BoolEntry{ID: pos, Val: x})
		}
		return out
	case rankStrCol:
		out := make(rcpg.SparseStr, 0, len(col.vals))
		for pos, v := range col.Entries() {
			x, _ := v.StrID()
			out = append(out, rcpg.StrEntry{ID: pos, Val: x})
		}
		return out
	}
	return nil
}

func dataToColumn(data rcpg.ColumnData) Column {
	switch d := data.(type) {
	case rcpg.DenseI64:
		return narrowI64Column(d)
	case rcpg.DenseF64:
		return denseF64Col(d)
	case rcpg.DenseBool:
		return denseBoolCol{bits: bitset.FromPackedLSB(d.Bits, int(d.Len))}
	case rcpg.DenseStr:
		return denseStrCol(d)
	case rcpg.SparseI64:
		col := sparseI64Col{ids: make([]uint32, len(d)), vals: make([]int64, len(d))}
		for i, e := range d {
			col.ids[i], col.vals[i] = e.ID, e.Val
		}
		if nv, ok := narrowI64Vals(col.vals); ok {
			return sparseI64NarrowCol{ids: col.ids, narrowI64Vec: nv}
		}
		return col
	case rcpg.SparseF64:
		col := sparseF64Col{ids: make([]uint32, len(d)), vals: make([]float64, len(d))}
		for i, e := range d {
			col.ids[i], col.vals[i] = e.ID, e.Val
		}
		return col
	case rcpg.SparseBool:
		col := sparseBoolCol{ids: make([]uint32, len(d)), vals: make([]bool, len(d))}
		for i, e := range d {
			col.ids[i], col.vals[i] = e.ID, e.Val
		}
		return col
	case rcpg.SparseStr:
		col := sparseStrCol{ids: make([]uint32, len(d)), vals: make([]uint32, len(d))}
		for i, e := range d {
			col.ids[i], col.vals[i] = e.ID, e.Val
		}
		return col
	}
	return nil
}

// ToGraphSection converts this snapshot to the plain on-disk data model.
func (g *Snapshot) ToGraphSection() *rcpg.GraphSection {
	return g.graphSectionWith(true, true)
}

// graphSectionWith builds the on-disk model, converting only the requested
// column sets so topology-only writes copy no property data. Index and
// column orders are sorted by atom, as the format requires.
func (g *Snapshot) graphSectionWith(nodeCols, relCols bool) *rcpg.GraphSection {
	labelIndex := make([]rcpg.BitmapEntry, 0, len(g.labelIndex))
	for l, set := range g.labelIndex {
		labelIndex = append(labelIndex, rcpg.BitmapEntry{Atom: l.ID(), Bitmap: set.Bitmap()})
	}
	sort.Slice(labelIndex, func(i, j int) bool { return labelIndex[i].Atom < labelIndex[j].Atom })

	typeIndex := make([]rcpg.BitmapEntry, 0, len(g.typeIndex))
	for t, set := range g.typeIndex {
		typeIndex = append(typeIndex, rcpg.BitmapEntry{Atom: t.ID(), Bitmap: set.Bitmap()})
	}
	sort.Slice(typeIndex, func(i, j int) bool { return typeIndex[i].Atom < typeIndex[j].Atom })

	convert := func(cols map[PropertyKey]Column, enabled bool) []rcpg.Column {
		if !enabled {
			return nil
		}
		out := make([]rcpg.Column, 0, len(cols))
		for key, col := range cols {
			out = append(out, rcpg.Column{Key: key, Data: columnToData(col)})
		}
		sort.Slice(out, func(i, j int) bool { return out[i].Key < out[j].Key })
		return out
	}

	outTypes := make([]uint32, len(g.outTypes))
	for i, t := range g.outTypes {
		outTypes[i] = t.ID()
	}
	inTypes := make([]uint32, len(g.inTypes))
	for i, t := range g.inTypes {
		inTypes[i] = t.ID()
	}

	// The existence section is emitted only for a genuinely sparse id
	// space: dense graphs (every loader) keep byte-identical files, and
	// absent-means-dense is exactly correct for them on read.
	var existence *roaring.Bitmap
	if g.existence != nil && uint32(g.existence.Len()) < g.CSRIDSpace() {
		existence = g.existence.Bitmap()
	}

	return &rcpg.GraphSection{
		NNodes:      g.nNodes,
		NRels:       g.nRels,
		OutOffsets:  g.outOffsets,
		OutNbrs:     g.outNbrs,
		OutTypes:    outTypes,
		InOffsets:   g.inOffsets,
		InNbrs:      g.inNbrs,
		InTypes:     inTypes,
		LabelIndex:  labelIndex,
		TypeIndex:   typeIndex,
		NodeColumns: convert(g.columns, nodeCols),
		RelColumns:  convert(g.relColumns, relCols),
		Version:     g.version,
		Atoms:       g.atoms.Strings(),
		Existence:   existence,
	}
}

// FromGraphSection builds a snapshot from the on-disk data model. The
// section's slices and bitmaps are taken over; the caller must not reuse
// them.
func FromGraphSection(section *rcpg.GraphSection) *Snapshot {
	g := newSnapshot()
	g.nNodes = section.NNodes
	g.nRels = section.NRels
	g.outOffsets = section.OutOffsets
	g.outNbrs = section.OutNbrs
	g.inOffsets = section.InOffsets
	g.inNbrs = section.InNbrs

	g.outTypes = make([]RelType, len(section.OutTypes))
	for i, t := range section.OutTypes {
		g.outTypes[i] = RelType(t)
	}
	g.inTypes = make([]RelType, len(section.InTypes))
	for i, t := range section.InTypes {
		g.inTypes[i] = RelType(t)
	}

	if section.Existence != nil {
		g.existence = nodeset.FromBitmap(section.Existence)
	}
	for _, e := range section.LabelIndex {
		g.labelIndex[Label(e.Atom)] = nodeset.FromBitmap(e.Bitmap)
	}
	for _, e := range section.TypeIndex {
		g.typeIndex[RelType(e.Atom)] = nodeset.FromBitmap(e.Bitmap)
	}
	for _, c := range section.NodeColumns {
		g.columns[c.Key] = dataToColumn(c.Data)
	}
	for _, c := range section.RelColumns {
		g.relColumns[c.Key] = dataToColumn(c.Data)
	}
	g.version = section.Version
	g.atoms = NewAtoms(section.Atoms)

	// The format does not store the incoming -> outgoing position map; it is
	// rebuilt lazily on the first incoming rel-property read (getInToOut), so
	// a loaded-but-unqueried graph never pays its footprint. Record only
	// whether rel columns exist, the authority for whether it is ever needed.
	g.hasRelProps = len(g.relColumns) > 0
	return g
}

// WriteRCPG serializes this snapshot to RCPG bytes including property
// columns (see rustychickpeas-format's FORMAT.md for the layout).
func (g *Snapshot) WriteRCPG(w io.Writer) error {
	return g.WriteRCPGWith(w, rcpg.DefaultWriteOptions())
}

// WriteRCPGWith serializes with optional sections per opts;
// TopologyOnlyWriteOptions produces a lean traversal-only file (no property
// columns converted or written).
func (g *Snapshot) WriteRCPGWith(w io.Writer, opts rcpg.WriteOptions) error {
	return rcpg.WriteWith(g.graphSectionWith(opts.NodeColumns, opts.RelColumns), w, opts)
}

// ReadRCPG reads a snapshot from RCPG bytes. Lazy indexes rebuild on first
// use.
func ReadRCPG(b []byte) (*Snapshot, error) {
	section, err := rcpg.Parse(b)
	if err != nil {
		return nil, err
	}
	return FromGraphSection(section), nil
}

// WriteRCPGFile serializes to an RCPG file on disk.
func (g *Snapshot) WriteRCPGFile(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	w := bufio.NewWriter(f)
	if err := g.WriteRCPG(w); err != nil {
		f.Close()
		return err
	}
	if err := w.Flush(); err != nil {
		f.Close()
		return err
	}
	return f.Close()
}

// ReadRCPGFile reads a snapshot from an RCPG file on disk.
func ReadRCPGFile(path string) (*Snapshot, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	return ReadRCPG(b)
}

The pages are generated with Golds v0.8.7. (GOOS=linux GOARCH=amd64)
Golds is a Go 101 project developed by Tapir Liu.
PR and bug reports are welcome and can be submitted to the issue list.
Please follow @zigo_101 (reachable from the left QR code) to get the latest news of Golds.