protocol.mx raw
1 package iskra
2
3 import "io"
4
5 // Op codes for the iskra wire protocol.
6 const (
7 OpBigramWeight uint8 = 0x01
8 OpBigramWeightRelaxed uint8 = 0x02
9 OpIngestBigram uint8 = 0x03
10 OpIngestText uint8 = 0x04
11 OpWalkStep uint8 = 0x05
12 OpTranslate uint8 = 0x06
13 OpLookup uint8 = 0x07
14 OpInsert uint8 = 0x08
15 OpBeamRun uint8 = 0x09
16 OpFlush uint8 = 0x0A
17 OpIngestRef uint8 = 0x0B
18 OpGetRefs uint8 = 0x0C
19 OpWalkOutward uint8 = 0x0D
20 )
21
22 // Response status codes.
23 const (
24 StatusOK uint8 = 0
25 StatusNotFound uint8 = 1
26 StatusError uint8 = 2
27 )
28
29 // Frame layout: [length:4 LE][op:1][payload:N]
30 // Response layout: [length:4 LE][status:1][payload:N]
31
32 // WriteFrame writes a framed message (length-prefixed) to w.
33 func WriteFrame(w io.Writer, data []byte) error {
34 var hdr [4]byte
35 n := uint32(len(data))
36 hdr[0] = byte(n)
37 hdr[1] = byte(n >> 8)
38 hdr[2] = byte(n >> 16)
39 hdr[3] = byte(n >> 24)
40 if _, err := w.Write(hdr[:]); err != nil {
41 return err
42 }
43 _, err := w.Write(data)
44 return err
45 }
46
47 // ReadFrame reads a length-prefixed frame from r.
48 // Returns nil on EOF or error.
49 func ReadFrame(r io.Reader) ([]byte, error) {
50 var hdr [4]byte
51 if _, err := io.ReadFull(r, hdr[:]); err != nil {
52 return nil, err
53 }
54 n := uint32(hdr[0]) | uint32(hdr[1])<<8 | uint32(hdr[2])<<16 | uint32(hdr[3])<<24
55 if n == 0 {
56 return []byte{}, nil
57 }
58 if n > 4<<20 {
59 return nil, io.ErrUnexpectedEOF
60 }
61 buf := []byte{:int32(n)}
62 if _, err := io.ReadFull(r, buf); err != nil {
63 return nil, err
64 }
65 return buf, nil
66 }
67
68 // EncodeString appends a length-prefixed string (uint16 LE + bytes).
69 func EncodeString(buf []byte, s string) []byte {
70 n := uint16(len(s))
71 buf = append(buf, byte(n), byte(n>>8))
72 buf = append(buf, []byte(s)...)
73 return buf
74 }
75
76 // DecodeString reads a length-prefixed string from buf at offset.
77 // Returns (string, bytesConsumed). Returns ("", 0) on underflow.
78 func DecodeString(buf []byte, off int32) (string, int32) {
79 if off+2 > len(buf) {
80 return "", 0
81 }
82 n := int32(uint16(buf[off]) | uint16(buf[off+1])<<8)
83 off += 2
84 if off+n > len(buf) {
85 return "", 0
86 }
87 return string(buf[off : off+n]), 2 + n
88 }
89
90 // EncodeUint32 appends a uint32 LE.
91 func EncodeUint32(buf []byte, v uint32) []byte {
92 return append(buf, byte(v), byte(v>>8), byte(v>>16), byte(v>>24))
93 }
94
95 // DecodeUint32 reads a uint32 LE from buf at offset.
96 func DecodeUint32(buf []byte, off int32) (uint32, int32) {
97 if off+4 > len(buf) {
98 return 0, 0
99 }
100 v := uint32(buf[off]) | uint32(buf[off+1])<<8 | uint32(buf[off+2])<<16 | uint32(buf[off+3])<<24
101 return v, 4
102 }
103
104 // EncodeUint64 appends a uint64 LE.
105 func EncodeUint64(buf []byte, v uint64) []byte {
106 return append(buf, byte(v), byte(v>>8), byte(v>>16), byte(v>>24),
107 byte(v>>32), byte(v>>40), byte(v>>48), byte(v>>56))
108 }
109
110 // DecodeUint64 reads a uint64 LE from buf at offset.
111 func DecodeUint64(buf []byte, off int32) (uint64, int32) {
112 if off+8 > len(buf) {
113 return 0, 0
114 }
115 v := uint64(buf[off]) | uint64(buf[off+1])<<8 | uint64(buf[off+2])<<16 | uint64(buf[off+3])<<24 |
116 uint64(buf[off+4])<<32 | uint64(buf[off+5])<<40 | uint64(buf[off+6])<<48 | uint64(buf[off+7])<<56
117 return v, 8
118 }
119
120 // Request encoders
121
122 func EncodeBigramWeightReq(domain uint8, coord uint64, prev, next string) []byte {
123 buf := []byte{:0:3 + 8 + len(prev) + len(next)}
124 buf = append(buf, OpBigramWeight, domain)
125 buf = EncodeUint64(buf, coord)
126 buf = EncodeString(buf, prev)
127 buf = EncodeString(buf, next)
128 return buf
129 }
130
131 func EncodeBigramWeightRelaxedReq(domain uint8, coord uint64, prev, next string) []byte {
132 buf := []byte{:0:3 + 8 + len(prev) + len(next)}
133 buf = append(buf, OpBigramWeightRelaxed, domain)
134 buf = EncodeUint64(buf, coord)
135 buf = EncodeString(buf, prev)
136 buf = EncodeString(buf, next)
137 return buf
138 }
139
140 func EncodeIngestBigramReq(domain uint8, coord uint64, prev, next string) []byte {
141 buf := []byte{:0:3 + 8 + len(prev) + len(next)}
142 buf = append(buf, OpIngestBigram, domain)
143 buf = EncodeUint64(buf, coord)
144 buf = EncodeString(buf, prev)
145 buf = EncodeString(buf, next)
146 return buf
147 }
148
149 func EncodeIngestTextReq(domain uint8, coord uint64, tokens []string) []byte {
150 total := 2 + 8 + 2
151 for _, tok := range tokens {
152 total += 2 + len(tok)
153 }
154 buf := []byte{:0:total}
155 buf = append(buf, OpIngestText, domain)
156 buf = EncodeUint64(buf, coord)
157 buf = append(buf, byte(len(tokens)), byte(len(tokens)>>8))
158 for _, tok := range tokens {
159 buf = EncodeString(buf, tok)
160 }
161 return buf
162 }
163
164 func EncodeWalkStepReq(domain uint8, coord uint64, word string, maxCandidates uint16) []byte {
165 buf := []byte{:0:2 + 8 + 2 + len(word) + 2}
166 buf = append(buf, OpWalkStep, domain)
167 buf = EncodeUint64(buf, coord)
168 buf = EncodeString(buf, word)
169 buf = append(buf, byte(maxCandidates), byte(maxCandidates>>8))
170 return buf
171 }
172
173 func EncodeTranslateReq(srcDomain, dstDomain uint8, coord uint64, word string) []byte {
174 buf := []byte{:0:3 + 8 + 2 + len(word)}
175 buf = append(buf, OpTranslate, srcDomain, dstDomain)
176 buf = EncodeUint64(buf, coord)
177 buf = EncodeString(buf, word)
178 return buf
179 }
180
181 func EncodeBeamRunReq(domain uint8, coord uint64, word string, width, maxSteps uint16, kind WalkKind) []byte {
182 buf := []byte{:0:2 + 8 + 2 + len(word) + 5}
183 buf = append(buf, OpBeamRun, domain)
184 buf = EncodeUint64(buf, coord)
185 buf = EncodeString(buf, word)
186 buf = append(buf, byte(width), byte(width>>8))
187 buf = append(buf, byte(maxSteps), byte(maxSteps>>8))
188 buf = append(buf, byte(kind))
189 return buf
190 }
191
192 func EncodeFlushReq() []byte {
193 return []byte{OpFlush}
194 }
195
196 func EncodeIngestRefReq(domain uint8, coord uint64, src, target string, kind uint8) []byte {
197 buf := []byte{:0:3 + 8 + len(src) + len(target) + 4}
198 buf = append(buf, OpIngestRef, domain)
199 buf = EncodeUint64(buf, coord)
200 buf = EncodeString(buf, src)
201 buf = EncodeString(buf, target)
202 buf = append(buf, kind)
203 return buf
204 }
205
206 func EncodeGetRefsReq(domain uint8, src string) []byte {
207 buf := []byte{:0:2 + 2 + len(src)}
208 buf = append(buf, OpGetRefs, domain)
209 buf = EncodeString(buf, src)
210 return buf
211 }
212
213 func EncodeWalkOutwardReq(domain uint8, src string, maxDepth uint8) []byte {
214 buf := []byte{:0:3 + 2 + len(src)}
215 buf = append(buf, OpWalkOutward, domain)
216 buf = EncodeString(buf, src)
217 buf = append(buf, maxDepth)
218 return buf
219 }
220
221 func EncodeRefsResponse(refs []RefEntry) []byte {
222 buf := []byte{:0:3 + len(refs)*16}
223 buf = append(buf, StatusOK)
224 buf = append(buf, byte(len(refs)), byte(len(refs)>>8))
225 for _, r := range refs {
226 buf = EncodeUint32(buf, r.RecIdx)
227 buf = EncodeString(buf, r.Target)
228 buf = append(buf, r.Kind)
229 buf = EncodeUint32(buf, r.Weight)
230 }
231 return buf
232 }
233
234 // Response encoders
235
236 func EncodeResponse(status uint8, payload []byte) []byte {
237 buf := []byte{:0:1 + len(payload)}
238 buf = append(buf, status)
239 buf = append(buf, payload...)
240 return buf
241 }
242
243 func EncodeUint32Response(v uint32) []byte {
244 buf := []byte{StatusOK, 0, 0, 0, 0}
245 buf[1] = byte(v)
246 buf[2] = byte(v >> 8)
247 buf[3] = byte(v >> 16)
248 buf[4] = byte(v >> 24)
249 return buf
250 }
251
252 func EncodeWeightRelaxedResponse(w uint32, coord uint64) []byte {
253 buf := []byte{:0:13}
254 buf = append(buf, StatusOK)
255 buf = EncodeUint32(buf, w)
256 buf = EncodeUint64(buf, coord)
257 return buf
258 }
259
260 func EncodeStringResponse(s string) []byte {
261 buf := []byte{:0:3 + len(s)}
262 buf = append(buf, StatusOK)
263 buf = EncodeString(buf, s)
264 return buf
265 }
266
267 func EncodeWalkCandidatesResponse(candidates []WalkCandidate) []byte {
268 buf := []byte{:0:3 + len(candidates)*32}
269 buf = append(buf, StatusOK)
270 buf = append(buf, byte(len(candidates)), byte(len(candidates)>>8))
271 for _, c := range candidates {
272 buf = append(buf, c.State.Domain)
273 buf = EncodeUint64(buf, c.State.Coord)
274 buf = EncodeString(buf, c.State.Word)
275 buf = EncodeUint32(buf, c.State.Score)
276 buf = append(buf, c.State.Depth)
277 buf = EncodeUint32(buf, c.Weight)
278 }
279 return buf
280 }
281
282 func EncodeBeamResponse(s WalkState) []byte {
283 buf := []byte{:0:16 + len(s.Word)}
284 buf = append(buf, StatusOK)
285 buf = append(buf, s.Domain, s.Depth)
286 buf = EncodeUint64(buf, s.Coord)
287 buf = EncodeUint32(buf, s.Score)
288 buf = EncodeString(buf, s.Word)
289 return buf
290 }
291