io.mx raw

   1  package lattice
   2  
   3  import (
   4  	"encoding/binary"
   5  	"fmt"
   6  	"io"
   7  	"os"
   8  )
   9  
  10  var le = binary.LittleEndian
  11  
  12  // nodeSize is the on-disk size of a C=8 Node.
  13  // 8×16B keys + 8×4B RecPtrs + 9×4B Children + 4B (Mult, Flags, pad) = 200 bytes.
  14  const nodeSize = C*16 + C*4 + (C+1)*4 + 4
  15  
  16  // headerSize for version 4: magic(4) + version(2) + nCount(4) + rCount(4) +
  17  // pendCount(4) + 8×roots(32) + freeHead(4) + freeCount(4) + nodeCap(4) + recCap(4) = 66 → pad to 72.
  18  const headerSize = 72
  19  
  20  type iskrHeader struct {
  21  	nCount, rCount, pendCount uint32
  22  	roots                     [C]uint32
  23  	freeHead                  uint32
  24  	freeCount                 uint32
  25  	nodeCap                   uint32
  26  	recCap                    uint32
  27  }
  28  
  29  func writeHeaderBuf(t *Tree, nodeCap, recCap uint32) [headerSize]byte {
  30  	var buf [headerSize]byte
  31  	copy(buf[0:4], "ISKR")
  32  	le.PutUint16(buf[4:6], 4) // version 4 = C=8
  33  	le.PutUint32(buf[6:10], t.nCount)
  34  	le.PutUint32(buf[10:14], t.rCount)
  35  	le.PutUint32(buf[14:18], uint32(len(t.Pending)))
  36  	for i := 0; i < C; i++ {
  37  		le.PutUint32(buf[18+i*4:], t.Roots[i])
  38  	}
  39  	// offset 18 + 8*4 = 50
  40  	le.PutUint32(buf[50:54], t.FreeHead)
  41  	le.PutUint32(buf[54:58], t.FreeCount)
  42  	le.PutUint32(buf[58:62], nodeCap)
  43  	le.PutUint32(buf[62:66], recCap)
  44  	// bytes 66-71 = padding (zero)
  45  	return buf
  46  }
  47  
  48  func readHeaderBuf(buf [headerSize]byte) (iskrHeader, error) {
  49  	if string(buf[0:4]) != "ISKR" {
  50  		return iskrHeader{}, fmt.Errorf("bad magic: %q", buf[0:4])
  51  	}
  52  	v := le.Uint16(buf[4:6])
  53  	if v != 4 {
  54  		return iskrHeader{}, fmt.Errorf("unsupported version %d (want 4)", v)
  55  	}
  56  	hdr := iskrHeader{
  57  		nCount:    le.Uint32(buf[6:10]),
  58  		rCount:    le.Uint32(buf[10:14]),
  59  		pendCount: le.Uint32(buf[14:18]),
  60  		freeHead:  le.Uint32(buf[50:54]),
  61  		freeCount: le.Uint32(buf[54:58]),
  62  		nodeCap:   le.Uint32(buf[58:62]),
  63  		recCap:    le.Uint32(buf[62:66]),
  64  	}
  65  	for i := 0; i < C; i++ {
  66  		hdr.roots[i] = le.Uint32(buf[18+i*4:])
  67  	}
  68  	return hdr, nil
  69  }
  70  
  71  func writeHeader(w io.Writer, t *Tree) error {
  72  	buf := writeHeaderBuf(t, t.nCount, t.rCount)
  73  	_, err := w.Write(buf[:])
  74  	return err
  75  }
  76  
  77  func readHeader(r io.Reader) (iskrHeader, error) {
  78  	var buf [headerSize]byte
  79  	if _, err := io.ReadFull(r, buf[:]); err != nil {
  80  		return iskrHeader{}, fmt.Errorf("read header: %w", err)
  81  	}
  82  	return readHeaderBuf(buf)
  83  }
  84  
  85  // Node is now nodeSize bytes: 8×16B keys + 8×4B RecPtrs + 9×4B Children + 4B flags/pad.
  86  func writeNodeBuf(n *Node) [nodeSize]byte {
  87  	var buf [nodeSize]byte
  88  	for i := 0; i < C; i++ {
  89  		le.PutUint64(buf[i*16:], n.Keys[i][0])
  90  		le.PutUint64(buf[i*16+8:], n.Keys[i][1])
  91  	}
  92  	// RecPtrs at offset C*16 = 128
  93  	for i := 0; i < C; i++ {
  94  		le.PutUint32(buf[128+i*4:], n.RecPtrs[i])
  95  	}
  96  	// Children at offset 128 + C*4 = 160
  97  	for i := 0; i < C+1; i++ {
  98  		le.PutUint32(buf[160+i*4:], n.Children[i])
  99  	}
 100  	// Mult and Flags at offset 160 + (C+1)*4 = 196
 101  	buf[196] = n.Mult
 102  	buf[197] = n.Flags
 103  	return buf
 104  }
 105  
 106  func readNodeBuf(buf [nodeSize]byte) Node {
 107  	var n Node
 108  	for i := 0; i < C; i++ {
 109  		n.Keys[i][0] = le.Uint64(buf[i*16:])
 110  		n.Keys[i][1] = le.Uint64(buf[i*16+8:])
 111  	}
 112  	for i := 0; i < C; i++ {
 113  		n.RecPtrs[i] = le.Uint32(buf[128+i*4:])
 114  	}
 115  	for i := 0; i < C+1; i++ {
 116  		n.Children[i] = le.Uint32(buf[160+i*4:])
 117  	}
 118  	n.Mult = buf[196]
 119  	n.Flags = buf[197]
 120  	return n
 121  }
 122  
 123  func writeRecBuf(r *Record) [48]byte {
 124  	var buf [48]byte
 125  	le.PutUint32(buf[0:], r.DataFile)
 126  	le.PutUint32(buf[4:], r.DataOff)
 127  	le.PutUint32(buf[8:], r.DataLen)
 128  	le.PutUint32(buf[12:], r.Link[0])
 129  	le.PutUint32(buf[16:], r.Link[1])
 130  	buf[20] = uint8(r.Branch)
 131  	copy(buf[24:], r.Inline[:])
 132  	return buf
 133  }
 134  
 135  func readRecBuf(buf [48]byte) Record {
 136  	var r Record
 137  	r.DataFile = le.Uint32(buf[0:])
 138  	r.DataOff = le.Uint32(buf[4:])
 139  	r.DataLen = le.Uint32(buf[8:])
 140  	r.Link[0] = le.Uint32(buf[12:])
 141  	r.Link[1] = le.Uint32(buf[16:])
 142  	r.Branch = buf[20]
 143  	copy(r.Inline[:], buf[24:])
 144  	return r
 145  }
 146  
 147  func writeNode(w io.Writer, n *Node) error {
 148  	buf := writeNodeBuf(n)
 149  	_, err := w.Write(buf[:])
 150  	return err
 151  }
 152  
 153  func readNode(r io.Reader) (Node, error) {
 154  	var buf [nodeSize]byte
 155  	if _, err := io.ReadFull(r, buf[:]); err != nil {
 156  		return Node{}, err
 157  	}
 158  	return readNodeBuf(buf), nil
 159  }
 160  
 161  func writeRecord(w io.Writer, r *Record) error {
 162  	buf := writeRecBuf(r)
 163  	_, err := w.Write(buf[:])
 164  	return err
 165  }
 166  
 167  func readRecord(r io.Reader) (Record, error) {
 168  	var buf [48]byte
 169  	if _, err := io.ReadFull(r, buf[:]); err != nil {
 170  		return Record{}, err
 171  	}
 172  	return readRecBuf(buf), nil
 173  }
 174  
 175  func Save(w io.Writer, t *Tree) error {
 176  	if err := writeHeader(w, t); err != nil {
 177  		return fmt.Errorf("header: %w", err)
 178  	}
 179  	for i := uint32(0); i < t.nCount; i++ {
 180  		if err := writeNode(w, &t.Nodes[i]); err != nil {
 181  			return fmt.Errorf("node %d: %w", i, err)
 182  		}
 183  	}
 184  	for i := uint32(0); i < t.rCount; i++ {
 185  		if err := writeRecord(w, &t.Records[i]); err != nil {
 186  			return fmt.Errorf("record %d: %w", i, err)
 187  		}
 188  	}
 189  	for i := range t.Pending {
 190  		if err := writePending(w, &t.Pending[i]); err != nil {
 191  			return fmt.Errorf("pending %d: %w", i, err)
 192  		}
 193  	}
 194  	return nil
 195  }
 196  
 197  func Load(r io.Reader) (*Tree, error) {
 198  	hdr, err := readHeader(r)
 199  	if err != nil {
 200  		return nil, err
 201  	}
 202  	t := &Tree{
 203  		Nodes:     []Node{:0:int(hdr.nCount)},
 204  		Records:   []Record{:0:int(hdr.rCount)},
 205  		Roots:     hdr.roots,
 206  		FreeHead:  hdr.freeHead,
 207  		FreeCount: hdr.freeCount,
 208  		nCount:    hdr.nCount,
 209  		rCount:    hdr.rCount,
 210  		RecKey:    map[uint32]Key{},
 211  	}
 212  	for i := uint32(0); i < hdr.nCount; i++ {
 213  		n, err := readNode(r)
 214  		if err != nil {
 215  			return nil, fmt.Errorf("node %d: %w", i, err)
 216  		}
 217  		t.Nodes = append(t.Nodes, n)
 218  	}
 219  	for i := uint32(0); i < hdr.rCount; i++ {
 220  		rec, err := readRecord(r)
 221  		if err != nil {
 222  			return nil, fmt.Errorf("record %d: %w", i, err)
 223  		}
 224  		t.Records = append(t.Records, rec)
 225  	}
 226  	if hdr.pendCount > 0 {
 227  		t.Pending = []PendingExp{:0:int(hdr.pendCount)}
 228  		for i := uint32(0); i < hdr.pendCount; i++ {
 229  			p, err := readPending(r)
 230  			if err != nil {
 231  				return nil, fmt.Errorf("pending %d: %w", i, err)
 232  			}
 233  			t.Pending = append(t.Pending, p)
 234  		}
 235  	}
 236  	t.rebuildRecKey()
 237  	return t, nil
 238  }
 239  
 240  func SaveFile(path string, t *Tree) error {
 241  	f, err := os.Create(path)
 242  	if err != nil {
 243  		return err
 244  	}
 245  	defer f.Close()
 246  	return Save(f, t)
 247  }
 248  
 249  func LoadFile(path string) (*Tree, error) {
 250  	f, err := os.Open(path)
 251  	if err != nil {
 252  		return nil, err
 253  	}
 254  	defer f.Close()
 255  	return Load(f)
 256  }
 257  
 258  func writePending(w io.Writer, p *PendingExp) error {
 259  	var buf [4]byte
 260  	buf[0] = uint8(p.Branch)
 261  	le.PutUint16(buf[1:3], p.Depth)
 262  	buf[3] = p.Factor
 263  	_, err := w.Write(buf[:])
 264  	return err
 265  }
 266  
 267  func readPending(r io.Reader) (PendingExp, error) {
 268  	var buf [4]byte
 269  	var p PendingExp
 270  	if _, err := io.ReadFull(r, buf[:]); err != nil {
 271  		return p, err
 272  	}
 273  	p.Branch = Branch(buf[0])
 274  	p.Depth = le.Uint16(buf[1:3])
 275  	p.Factor = buf[3]
 276  	return p, nil
 277  }
 278