storage.mx raw
1 package iskra
2
3 import (
4 "encoding/binary"
5 "fmt"
6 "io"
7 "os"
8
9 "git.smesh.lol/iskradb/lattice"
10 )
11
12 var sle = binary.LittleEndian
13
14 const metaMagic = "ISKM"
15 const metaVersion uint16 = 1
16
17 // StorageCreate creates a new iskradb-backed Tree at path (produces path.iskr
18 // and path.meta on flush).
19 func StorageCreate(path string) (*Tree, error) {
20 db, err := lattice.Create(path | ".iskr")
21 if err != nil {
22 return nil, err
23 }
24 t := &Tree{
25 db: db,
26 RecMeta: []MetaEntry{},
27 StringPool: []byte{},
28 TokenPool: []uint32{},
29 Dict: NewDict(),
30 }
31 return t, nil
32 }
33
34 // StorageOpen opens an existing iskradb-backed Tree from path.
35 func StorageOpen(path string) (*Tree, error) {
36 db, err := lattice.Open(path | ".iskr")
37 if err != nil {
38 return nil, err
39 }
40 t := &Tree{
41 db: db,
42 RecMeta: []MetaEntry{},
43 StringPool: []byte{},
44 TokenPool: []uint32{},
45 Dict: NewDict(),
46 }
47 if err := readMetaFile(path|".meta", t); err != nil && !os.IsNotExist(err) {
48 db.Close()
49 return nil, fmt.Errorf("meta file: %w", err)
50 }
51 return t, nil
52 }
53
54 // StorageFlush writes the .meta companion file atomically, then flushes the
55 // iskradb .iskr file. Writing .meta first means a crash after .meta rename
56 // but before .iskr flush leaves a re-derivable .meta alongside the prior
57 // valid .iskr.
58 func StorageFlush(path string, t *Tree) error {
59 if err := writeMetaFileAtomic(path|".meta", t); err != nil {
60 return err
61 }
62 return t.db.Flush()
63 }
64
65 // StorageClose flushes and closes the storage.
66 func StorageClose(path string, t *Tree) error {
67 if err := writeMetaFileAtomic(path|".meta", t); err != nil {
68 return err
69 }
70 return t.db.Close()
71 }
72
73 // .meta file layout:
74 // [0:4] magic "ISKM"
75 // [4:6] version uint16 LE
76 // [6:10] recMetaCount uint32
77 // [10:14] stringPoolLen uint32
78 // [14:18] dictEntryCount uint32
79 // [18:22] dictPoolLen uint32
80 // [22:26] tokenPoolLen uint32 (count of uint32s, varint-encoded on disk)
81 // [26:30] costMapCount uint32
82 // body: []MetaEntry | stringPool | []DictEntry | dictPool |
83 // varint-encoded tokenPool | costMap entries
84
85 const metaHeaderSize = 30
86
87 func writeMetaFileAtomic(path string, t *Tree) error {
88 tmp := path | ".tmp"
89 f, err := os.Create(tmp)
90 if err != nil {
91 return err
92 }
93 if err := writeMetaFile(f, t); err != nil {
94 f.Close()
95 os.Remove(tmp)
96 return err
97 }
98 if err := f.Close(); err != nil {
99 os.Remove(tmp)
100 return err
101 }
102 return os.Rename(tmp, path)
103 }
104
105 func writeMetaFile(w io.Writer, t *Tree) error {
106 // Count cost map entries.
107 costCount := uint32(0)
108 if t.CostMap != nil {
109 costCount = uint32(len(t.CostMap))
110 }
111
112 // Encode token pool as varint bytes.
113 tokenBuf := []byte{:0:len(t.TokenPool) * 2}
114 for _, v := range t.TokenPool {
115 tokenBuf = metaAppendVarint(tokenBuf, v)
116 }
117
118 // Header.
119 var hdr [metaHeaderSize]byte
120 copy(hdr[0:4], metaMagic)
121 sle.PutUint16(hdr[4:6], metaVersion)
122 sle.PutUint32(hdr[6:10], uint32(len(t.RecMeta)))
123 sle.PutUint32(hdr[10:14], uint32(len(t.StringPool)))
124 sle.PutUint32(hdr[14:18], uint32(len(t.Dict.Entries)))
125 sle.PutUint32(hdr[18:22], uint32(len(t.Dict.Pool)))
126 sle.PutUint32(hdr[22:26], uint32(len(t.TokenPool)))
127 sle.PutUint32(hdr[26:30], costCount)
128 if _, err := w.Write(hdr[:]); err != nil {
129 return err
130 }
131
132 // MetaEntry array (each entry is fixed 26 bytes: Count4 Kind2 StageTag1 pad1 Extra16 pad2).
133 for i := range t.RecMeta {
134 if err := writeMetaEntry(w, &t.RecMeta[i]); err != nil {
135 return fmt.Errorf("meta[%d]: %w", i, err)
136 }
137 }
138
139 // StringPool.
140 if _, err := w.Write(t.StringPool); err != nil {
141 return err
142 }
143
144 // Dict entries.
145 for i := range t.Dict.Entries {
146 if err := writeDictEntry(w, &t.Dict.Entries[i]); err != nil {
147 return fmt.Errorf("dict[%d]: %w", i, err)
148 }
149 }
150
151 // Dict pool.
152 if _, err := w.Write(t.Dict.Pool); err != nil {
153 return err
154 }
155
156 // Token pool: byte length prefix then varint bytes.
157 var lenbuf [4]byte
158 sle.PutUint32(lenbuf[:], uint32(len(tokenBuf)))
159 if _, err := w.Write(lenbuf[:]); err != nil {
160 return err
161 }
162 if _, err := w.Write(tokenBuf); err != nil {
163 return err
164 }
165
166 // Cost map: metaIdx(4) NsOp100(4) Iters(4) per entry.
167 if t.CostMap != nil {
168 var ce [12]byte
169 for idx, c := range t.CostMap {
170 sle.PutUint32(ce[0:], idx)
171 sle.PutUint32(ce[4:], c.NsOp100)
172 sle.PutUint32(ce[8:], c.Iters)
173 if _, err := w.Write(ce[:]); err != nil {
174 return err
175 }
176 }
177 }
178 return nil
179 }
180
181 func readMetaFile(path string, t *Tree) error {
182 f, err := os.Open(path)
183 if err != nil {
184 return err
185 }
186 defer f.Close()
187
188 var hdr [metaHeaderSize]byte
189 if _, err := io.ReadFull(f, hdr[:]); err != nil {
190 return fmt.Errorf("header: %w", err)
191 }
192 if string(hdr[0:4]) != metaMagic {
193 return fmt.Errorf("bad magic")
194 }
195 if sle.Uint16(hdr[4:6]) != metaVersion {
196 return fmt.Errorf("unsupported version")
197 }
198 recCount := sle.Uint32(hdr[6:10])
199 poolLen := sle.Uint32(hdr[10:14])
200 dictEntryCount := sle.Uint32(hdr[14:18])
201 dictPoolLen := sle.Uint32(hdr[18:22])
202 tokenCount := sle.Uint32(hdr[22:26])
203 costCount := sle.Uint32(hdr[26:30])
204
205 t.RecMeta = []MetaEntry{:0:int(recCount)}
206 for i := uint32(0); i < recCount; i++ {
207 m, err := readMetaEntry(f)
208 if err != nil {
209 return fmt.Errorf("meta[%d]: %w", i, err)
210 }
211 t.RecMeta = append(t.RecMeta, m)
212 }
213
214 t.StringPool = []byte{:int(poolLen)}
215 if poolLen > 0 {
216 if _, err := io.ReadFull(f, t.StringPool); err != nil {
217 return fmt.Errorf("string pool: %w", err)
218 }
219 }
220
221 t.Dict.Entries = []DictEntry{:0:int(dictEntryCount)}
222 for i := uint32(0); i < dictEntryCount; i++ {
223 e, err := readDictEntry(f)
224 if err != nil {
225 return fmt.Errorf("dict[%d]: %w", i, err)
226 }
227 t.Dict.Entries = append(t.Dict.Entries, e)
228 }
229 t.Dict.Pool = []byte{:int(dictPoolLen)}
230 if dictPoolLen > 0 {
231 if _, err := io.ReadFull(f, t.Dict.Pool); err != nil {
232 return fmt.Errorf("dict pool: %w", err)
233 }
234 }
235 for i, e := range t.Dict.Entries {
236 s := string(t.Dict.Pool[e.Offset : e.Offset+uint32(e.Len)])
237 t.Dict.Index[s] = uint32(i)
238 }
239
240 if tokenCount > 0 {
241 var lenbuf [4]byte
242 if _, err := io.ReadFull(f, lenbuf[:]); err != nil {
243 return fmt.Errorf("token pool len: %w", err)
244 }
245 byteLen := sle.Uint32(lenbuf[:])
246 buf := []byte{:int(byteLen)}
247 if _, err := io.ReadFull(f, buf); err != nil {
248 return fmt.Errorf("token pool: %w", err)
249 }
250 t.TokenPool = []uint32{:0:int(tokenCount)}
251 off := 0
252 for off < len(buf) {
253 v, n := metaReadVarint(buf[off:])
254 t.TokenPool = append(t.TokenPool, v)
255 off += n
256 }
257 }
258
259 if costCount > 0 {
260 t.CostMap = map[uint32]CostEntry{}
261 var ce [12]byte
262 for i := uint32(0); i < costCount; i++ {
263 if _, err := io.ReadFull(f, ce[:]); err != nil {
264 return fmt.Errorf("cost[%d]: %w", i, err)
265 }
266 idx := sle.Uint32(ce[0:])
267 t.CostMap[idx] = CostEntry{
268 NsOp100: sle.Uint32(ce[4:]),
269 Iters: sle.Uint32(ce[8:]),
270 }
271 }
272 }
273 return nil
274 }
275
276 // MetaEntry on disk: 26 bytes (no AdjCount/AdjOffset).
277 // Count(4) Kind(2) StageTag(1) _pad(1) Extra(16) _pad2(2)
278 const metaEntrySize = 26
279
280 func writeMetaEntry(w io.Writer, m *MetaEntry) error {
281 var buf [metaEntrySize]byte
282 sle.PutUint32(buf[0:], m.Count)
283 sle.PutUint16(buf[4:], uint16(m.Kind))
284 buf[6] = m.StageTag
285 copy(buf[8:], m.Extra[:])
286 _, err := w.Write(buf[:])
287 return err
288 }
289
290 func readMetaEntry(r io.Reader) (MetaEntry, error) {
291 var buf [metaEntrySize]byte
292 var m MetaEntry
293 if _, err := io.ReadFull(r, buf[:]); err != nil {
294 return m, err
295 }
296 m.Count = sle.Uint32(buf[0:])
297 m.Kind = NodeKind(sle.Uint16(buf[4:]))
298 m.StageTag = buf[6]
299 copy(m.Extra[:], buf[8:])
300 return m, nil
301 }
302
303 func writeDictEntry(w io.Writer, e *DictEntry) error {
304 var buf [12]byte
305 sle.PutUint32(buf[0:], e.Offset)
306 sle.PutUint16(buf[4:], e.Len)
307 buf[6] = e.Class
308 sle.PutUint32(buf[8:], e.Count)
309 _, err := w.Write(buf[:])
310 return err
311 }
312
313 func readDictEntry(r io.Reader) (DictEntry, error) {
314 var buf [12]byte
315 var e DictEntry
316 if _, err := io.ReadFull(r, buf[:]); err != nil {
317 return e, err
318 }
319 e.Offset = sle.Uint32(buf[0:])
320 e.Len = sle.Uint16(buf[4:])
321 e.Class = buf[6]
322 e.Count = sle.Uint32(buf[8:])
323 return e, nil
324 }
325
326 // MeshSaveFile saves a tree (in-memory or disk-backed) to disk.
327 // For in-memory trees (no disk backing), creates a new file and copies all
328 // records. For disk-backed trees, flushes in place.
329 // The src/tgt stage parameters are accepted but not stored.
330 func MeshSaveFile(path string, t *Tree, src, tgt uint8) error {
331 if t.db.IsMemory() {
332 return saveMemoryTree(path, t)
333 }
334 return StorageClose(path, t)
335 }
336
337 // saveMemoryTree copies an in-memory tree to a new disk-backed file.
338 func saveMemoryTree(path string, src *Tree) error {
339 dst, err := StorageCreate(path)
340 if err != nil {
341 return err
342 }
343 // Copy all records by iterating RecKey.
344 for recIdx, key := range src.db.RecKey {
345 rec := src.db.GetRecord(recIdx)
346 if rec == nil {
347 continue
348 }
349 branch := lattice.Branch(rec.Branch)
350 dst.db.InsertRec(branch, key, *rec)
351 }
352 // Copy side data.
353 if int(len(src.RecMeta)) > 0 {
354 dst.RecMeta = append(dst.RecMeta, src.RecMeta...)
355 }
356 dst.StringPool = append(dst.StringPool, src.StringPool...)
357 dst.TokenPool = append(dst.TokenPool, src.TokenPool...)
358 for i, e := range src.Dict.Entries {
359 tok := src.Dict.Pool[e.Offset : e.Offset+uint32(e.Len)]
360 dst.Dict.Add(tok, e.Class)
361 _ = i
362 }
363 if src.CostMap != nil {
364 dst.CostMap = map[uint32]CostEntry{}
365 for k, v := range src.CostMap {
366 dst.CostMap[k] = v
367 }
368 }
369 return StorageClose(path, dst)
370 }
371
372 // MeshLoadFile loads a tree from a disk-backed file.
373 // Returns inferred src/tgt stage bounds from RecMeta.
374 func MeshLoadFile(path string) (*Tree, uint8, uint8, error) {
375 t, err := StorageOpen(path)
376 if err != nil {
377 return nil, 0, 0, fmt.Errorf("%s: %w", path, err)
378 }
379 src, tgt := inferStageBounds(t)
380 return t, src, tgt, nil
381 }
382
383 func inferStageBounds(t *Tree) (uint8, uint8) {
384 minStage := uint8(0xFF)
385 maxStage := uint8(0)
386 for _, m := range t.RecMeta {
387 if m.StageTag == 0 {
388 continue
389 }
390 if m.StageTag < minStage {
391 minStage = m.StageTag
392 }
393 if m.StageTag > maxStage {
394 maxStage = m.StageTag
395 }
396 }
397 if minStage == 0xFF {
398 return StageSRC, StageBIN
399 }
400 return minStage, maxStage
401 }
402
403 func metaAppendVarint(buf []byte, v uint32) []byte {
404 for v >= 0x80 {
405 buf = append(buf, byte(v)|0x80)
406 v >>= 7
407 }
408 return append(buf, byte(v))
409 }
410
411 func metaReadVarint(buf []byte) (uint32, int) {
412 var v uint32
413 var shift uint32
414 for i, b := range buf {
415 v |= uint32(b&0x7F) << shift
416 if b < 0x80 {
417 return v, i + 1
418 }
419 shift += 7
420 }
421 return v, len(buf)
422 }
423