transfer.mx raw
1 package lattice
2
3 type ExportFilter func(nodeIdx uint32, n *Node, recIdx uint32, r *Record) bool
4
5 type SubLattice struct {
6 Tree *Tree
7 OriginKey Key
8 SrcBranch Branch
9 }
10
11 type keyRec struct {
12 key Key
13 recIdx uint32
14 branch Branch
15 }
16
17 func (t *Tree) collectBranch(rootIdx uint32, branch Branch, filter ExportFilter) []keyRec {
18 var out []keyRec
19 queue := []uint32{rootIdx}
20 for len(queue) > 0 {
21 idx := queue[0]
22 queue = queue[1:]
23 if idx == NullIndex {
24 continue
25 }
26 node := t.getNode(idx)
27 if node.IsDeleted() {
28 continue
29 }
30 k := node.KeyCount()
31 for i := 0; i < k; i++ {
32 ri := node.RecPtrs[i]
33 if ri == NullRec {
34 continue
35 }
36 if filter != nil {
37 rec := t.getRecord(ri)
38 if !filter(idx, node, ri, rec) {
39 continue
40 }
41 }
42 out = append(out, keyRec{key: node.Keys[i], recIdx: ri, branch: branch})
43 }
44 for i := 0; i <= k; i++ {
45 if node.Children[i] != NullIndex {
46 queue = append(queue, node.Children[i])
47 }
48 }
49 }
50 return out
51 }
52
53 func (t *Tree) extractPairs(pairs []keyRec) *SubLattice {
54 dst := NewTree(len(pairs) + 16)
55 remap := map[uint32]uint32{}
56 for _, p := range pairs {
57 rec := *t.getRecord(p.recIdx)
58 rec.Link[0] = NullRec
59 rec.Link[1] = NullRec
60 dst.Insert(p.branch, p.key, rec)
61 newRI := dst.LookupRecIdx(p.branch, p.key)
62 remap[p.recIdx] = newRI
63 }
64 for _, p := range pairs {
65 srcRec := t.getRecord(p.recIdx)
66 newRI, ok := remap[p.recIdx]
67 if !ok {
68 continue
69 }
70 dstRec := dst.getRecord(newRI)
71 for li := 0; li < 2; li++ {
72 if srcRec.Link[li] == NullRec {
73 continue
74 }
75 if mapped, ok := remap[srcRec.Link[li]]; ok {
76 dstRec.Link[li] = mapped
77 }
78 }
79 dstRec.Branch = srcRec.Branch
80 }
81 return &SubLattice{Tree: dst}
82 }
83
84 func (t *Tree) ExtractBranch(branch Branch) *SubLattice {
85 pairs := t.collectBranch(t.Roots[branch], branch, nil)
86 sl := t.extractPairs(pairs)
87 sl.SrcBranch = branch
88 return sl
89 }
90
91 func (t *Tree) ExtractSubtree(rootNodeIdx uint32, branch Branch, filter ExportFilter) *SubLattice {
92 pairs := t.collectBranch(rootNodeIdx, branch, filter)
93 sl := t.extractPairs(pairs)
94 sl.SrcBranch = branch
95 return sl
96 }
97
98 func (t *Tree) ExtractLinked(startRecIdx uint32) *SubLattice {
99 visited := map[uint32]bool{}
100 queue := []uint32{startRecIdx}
101 visited[startRecIdx] = true
102 for len(queue) > 0 {
103 ri := queue[0]
104 queue = queue[1:]
105 rec := t.getRecord(ri)
106 for li := 0; li < 2; li++ {
107 partner := rec.Link[li]
108 if partner == NullRec || visited[partner] {
109 continue
110 }
111 visited[partner] = true
112 queue = append(queue, partner)
113 }
114 }
115 var pairs []keyRec
116 for ri := range visited {
117 key, ok := t.RecKey[ri]
118 if !ok {
119 continue
120 }
121 rec := t.getRecord(ri)
122 pairs = append(pairs, keyRec{key: key, recIdx: ri, branch: Branch(rec.Branch)})
123 }
124 sl := t.extractPairs(pairs)
125 if startRecIdx != NullRec {
126 sl.OriginKey = t.RecKey[startRecIdx]
127 sl.SrcBranch = Branch(t.getRecord(startRecIdx).Branch)
128 }
129 return sl
130 }
131
132 func (t *Tree) Graft(sl *SubLattice) int {
133 var allPairs []keyRec
134 for b := 0; b < 3; b++ {
135 br := Branch(b)
136 pairs := sl.Tree.collectBranch(sl.Tree.Roots[br], br, nil)
137 allPairs = append(allPairs, pairs...)
138 }
139 remap := map[uint32]uint32{}
140 inserted := 0
141 for _, p := range allPairs {
142 existing := t.LookupRecIdx(p.branch, p.key)
143 if existing != NullRec {
144 remap[p.recIdx] = existing
145 continue
146 }
147 rec := *sl.Tree.getRecord(p.recIdx)
148 rec.Link[0] = NullRec
149 rec.Link[1] = NullRec
150 t.Insert(p.branch, p.key, rec)
151 newRI := t.LookupRecIdx(p.branch, p.key)
152 remap[p.recIdx] = newRI
153 inserted++
154 }
155 for _, p := range allPairs {
156 srcRec := sl.Tree.getRecord(p.recIdx)
157 newRI, ok := remap[p.recIdx]
158 if !ok {
159 continue
160 }
161 dstRec := t.getRecord(newRI)
162 for li := 0; li < 2; li++ {
163 if srcRec.Link[li] == NullRec {
164 continue
165 }
166 if mapped, ok := remap[srcRec.Link[li]]; ok {
167 dstRec.Link[li] = mapped
168 t.markRecDirty(newRI)
169 }
170 }
171 if srcRec.Branch != 0 || dstRec.Branch == 0 {
172 dstRec.Branch = srcRec.Branch
173 }
174 }
175 return inserted
176 }
177
178 func (sl *SubLattice) SaveFile(path string) error {
179 return SaveFile(path, sl.Tree)
180 }
181
182 func LoadSubLattice(path string) (*SubLattice, error) {
183 t, err := LoadFile(path)
184 if err != nil {
185 return nil, err
186 }
187 return &SubLattice{Tree: t}, nil
188 }
189