goenzyme.go raw
1 package enzyme
2
3 import (
4 "bytes"
5 "fmt"
6 "go/ast"
7 "go/parser"
8 "go/printer"
9 "go/token"
10 "io"
11 "regexp"
12 "strings"
13
14 "git.mleku.dev/mleku/dendrite/pkg/axiom"
15 )
16
17 // GoSource is an enzyme that decomposes Go source code into typed AST elements.
18 //
19 // Each emitted element carries two pieces of information:
20 // - Type tag: the AST node kind ("func", "assign", "return", "if", etc.)
21 // - Value: rendered source text, optionally prefixed with parent context
22 //
23 // Body-level statements (assign, return, if, for, switch, select, go, send)
24 // carry their parent function name separated by \x00:
25 //
26 // "main\x00x := foo()" — assignment inside main()
27 // "Run\x00return nil" — return inside Run()
28 //
29 // This allows the emitter to reconstruct function bodies by grouping
30 // elements that share a parent function.
31 type GoSource struct{}
32
33 // CanDigest returns true if the sample looks like Go source.
34 func (GoSource) CanDigest(sample []byte) bool {
35 for i := 0; i < len(sample)-8 && i < 512; i++ {
36 if string(sample[i:i+8]) == "package " {
37 return true
38 }
39 }
40 return false
41 }
42
43 // Digest parses Go source and emits typed elements for each declaration,
44 // statement, identifier, and literal. Body-level statements carry their
45 // rendered source text and parent function context.
46 func (GoSource) Digest(r io.Reader) <-chan axiom.Element {
47 ch := make(chan axiom.Element, 128)
48
49 go func() {
50 defer close(ch)
51
52 src, err := io.ReadAll(r)
53 if err != nil {
54 return
55 }
56
57 // Extract //go: directives before AST parsing (parser strips them).
58 emitDirectives(ch, src)
59
60 fset := token.NewFileSet()
61 file, err := parser.ParseFile(fset, "input.go", src, parser.ParseComments)
62 if err != nil {
63 return
64 }
65
66 // Package name.
67 ch <- newHexElement("package", file.Name.Name)
68
69 // Imports.
70 for _, imp := range file.Imports {
71 if imp.Path != nil {
72 ch <- newHexElement("import", imp.Path.Value)
73 }
74 }
75
76 // Process top-level declarations with scope tracking.
77 for _, decl := range file.Decls {
78 emitDecl(ch, fset, src, decl)
79 }
80 }()
81
82 return ch
83 }
84
85 // emitDecl processes a top-level declaration and emits elements.
86 func emitDecl(ch chan<- axiom.Element, fset *token.FileSet, src []byte, decl ast.Decl) {
87 switch d := decl.(type) {
88 case *ast.FuncDecl:
89 emitFuncDecl(ch, fset, src, d)
90 case *ast.GenDecl:
91 emitGenDecl(ch, fset, src, d)
92 }
93 }
94
95 // emitGenDecl processes type, const, and var declarations.
96 func emitGenDecl(ch chan<- axiom.Element, fset *token.FileSet, _ []byte, d *ast.GenDecl) {
97 for _, spec := range d.Specs {
98 switch s := spec.(type) {
99 case *ast.TypeSpec:
100 ch <- newHexElement("type", s.Name.Name)
101 ch <- newHexElement("ident:type-name", s.Name.Name)
102 switch s.Type.(type) {
103 case *ast.StructType:
104 ch <- newHexElement("struct", s.Name.Name)
105 case *ast.InterfaceType:
106 ch <- newHexElement("interface", s.Name.Name)
107 }
108 // Emit fields for struct types.
109 if st, ok := s.Type.(*ast.StructType); ok && st.Fields != nil {
110 for _, f := range st.Fields.List {
111 for _, name := range f.Names {
112 fieldVal := name.Name
113 if f.Type != nil {
114 fieldVal += " " + renderNode(fset, f.Type)
115 }
116 ch <- newHexElement("field", s.Name.Name+"\x00"+fieldVal)
117 ch <- newHexElement("ident:field-name", name.Name)
118 }
119 }
120 }
121 // Emit method signatures for interface types.
122 if it, ok := s.Type.(*ast.InterfaceType); ok && it.Methods != nil {
123 for _, m := range it.Methods.List {
124 for _, name := range m.Names {
125 sig := name.Name
126 if ft, ok := m.Type.(*ast.FuncType); ok {
127 sig += renderFuncSig(fset, ft)
128 }
129 ch <- newHexElement("method", sig)
130 ch <- newHexElement("ident:method-name", name.Name)
131 }
132 }
133 }
134 case *ast.ImportSpec:
135 // Already handled above.
136 case *ast.ValueSpec:
137 for _, name := range s.Names {
138 val := name.Name
139 if s.Type != nil {
140 val += " " + renderNode(fset, s.Type)
141 }
142 ch <- newHexElement("ident:var-name", val)
143 }
144 // Also emit a full var declaration for top-level reconstruction.
145 if d.Tok.String() == "var" || d.Tok.String() == "const" {
146 rendered := renderNode(fset, s)
147 ch <- newHexElement("var", d.Tok.String()+" "+rendered)
148 }
149 }
150 }
151 }
152
153 // emitFuncDecl processes a function declaration and its body.
154 func emitFuncDecl(ch chan<- axiom.Element, fset *token.FileSet, src []byte, fn *ast.FuncDecl) {
155 name := fn.Name.Name
156 if fn.Recv != nil {
157 // Method — emit with receiver type and variable name.
158 recv := ""
159 recvVar := ""
160 if len(fn.Recv.List) > 0 {
161 recv = renderNode(fset, fn.Recv.List[0].Type)
162 if len(fn.Recv.List[0].Names) > 0 {
163 recvVar = fn.Recv.List[0].Names[0].Name
164 }
165 }
166 ch <- newHexElement("method", recvVar+"\x00"+recv+"."+name+renderFuncSig(fset, fn.Type))
167 ch <- newHexElement("ident:method-name", name)
168 if recvVar != "" {
169 ch <- newHexElement("ident:receiver", recvVar)
170 }
171 } else {
172 ch <- newHexElement("func", name+renderFuncSig(fset, fn.Type))
173 ch <- newHexElement("ident:func-name", name)
174 }
175
176 // Emit param and result name idents.
177 emitFuncTypeIdents(ch, fn.Type)
178
179 // Emit function body statements with parent context.
180 if fn.Body != nil {
181 emitBlock(ch, fset, src, name, fn.Body)
182 }
183 }
184
185 // emitBlock processes a block statement, emitting each statement with
186 // its parent function context.
187 func emitBlock(ch chan<- axiom.Element, fset *token.FileSet, src []byte, parent string, block *ast.BlockStmt) {
188 if block == nil {
189 return
190 }
191 for _, stmt := range block.List {
192 emitStmt(ch, fset, src, parent, stmt)
193 }
194 }
195
196 // emitStmt processes a single statement, rendering it to source text
197 // and tagging it with its parent function and source line number.
198 func emitStmt(ch chan<- axiom.Element, fset *token.FileSet, src []byte, parent string, stmt ast.Stmt) {
199 // The value carries parent context and source position:
200 // "funcName\x00linenum\x00rendered_source"
201 // The line number enables the emitter to reconstruct original statement
202 // order when causal (define-use) analysis is ambiguous.
203 tag := stmtTag(stmt)
204 if tag == "" {
205 return
206 }
207
208 rendered := renderNode(fset, stmt)
209 lineNum := fset.Position(stmt.Pos()).Line
210 val := fmt.Sprintf("%s\x00%d\x00%s", parent, lineNum, rendered)
211
212 ch <- newHexElement(tag, val)
213
214 // Recurse into nested blocks so inner statements are also captured.
215 switch s := stmt.(type) {
216 case *ast.IfStmt:
217 emitBlock(ch, fset, src, parent, s.Body)
218 if s.Else != nil {
219 if elseBlock, ok := s.Else.(*ast.BlockStmt); ok {
220 emitBlock(ch, fset, src, parent, elseBlock)
221 } else if elseIf, ok := s.Else.(*ast.IfStmt); ok {
222 emitStmt(ch, fset, src, parent, elseIf)
223 }
224 }
225 case *ast.ForStmt:
226 emitBlock(ch, fset, src, parent, s.Body)
227 case *ast.RangeStmt:
228 emitBlock(ch, fset, src, parent, s.Body)
229 case *ast.SwitchStmt:
230 emitBlock(ch, fset, src, parent, s.Body)
231 case *ast.TypeSwitchStmt:
232 emitBlock(ch, fset, src, parent, s.Body)
233 case *ast.SelectStmt:
234 emitBlock(ch, fset, src, parent, s.Body)
235 case *ast.BlockStmt:
236 emitBlock(ch, fset, src, parent, s)
237 }
238
239 }
240
241 // stmtTag returns the element type tag for a statement node.
242 func stmtTag(stmt ast.Stmt) string {
243 switch stmt.(type) {
244 case *ast.AssignStmt:
245 return "assign"
246 case *ast.ReturnStmt:
247 return "return"
248 case *ast.IfStmt:
249 return "if"
250 case *ast.ForStmt, *ast.RangeStmt:
251 return "for"
252 case *ast.SwitchStmt, *ast.TypeSwitchStmt:
253 return "switch"
254 case *ast.SelectStmt:
255 return "select"
256 case *ast.GoStmt:
257 return "go"
258 case *ast.SendStmt:
259 return "send"
260 case *ast.ExprStmt:
261 return "expr"
262 case *ast.DeferStmt:
263 return "defer"
264 case *ast.DeclStmt:
265 return "decl"
266 case *ast.IncDecStmt:
267 return "assign" // treat i++ as assignment
268 case *ast.BranchStmt:
269 return "branch"
270 case *ast.CaseClause:
271 return "case"
272 case *ast.CommClause:
273 return "comm"
274 }
275 return ""
276 }
277
278 // renderNode renders an AST node back to Go source text.
279 func renderNode(fset *token.FileSet, node ast.Node) string {
280 var buf bytes.Buffer
281 cfg := printer.Config{Mode: printer.RawFormat, Tabwidth: 8}
282 if err := cfg.Fprint(&buf, fset, node); err != nil {
283 return ""
284 }
285 return strings.TrimSpace(buf.String())
286 }
287
288 // renderFuncSig renders a function type's parameter and result lists.
289 func renderFuncSig(fset *token.FileSet, ft *ast.FuncType) string {
290 if ft == nil {
291 return "()"
292 }
293 var buf bytes.Buffer
294 buf.WriteByte('(')
295 if ft.Params != nil {
296 for i, p := range ft.Params.List {
297 if i > 0 {
298 buf.WriteString(", ")
299 }
300 for j, name := range p.Names {
301 if j > 0 {
302 buf.WriteString(", ")
303 }
304 buf.WriteString(name.Name)
305 }
306 if p.Type != nil {
307 if len(p.Names) > 0 {
308 buf.WriteByte(' ')
309 }
310 buf.WriteString(renderNode(fset, p.Type))
311 }
312 }
313 }
314 buf.WriteByte(')')
315 if ft.Results != nil && len(ft.Results.List) > 0 {
316 buf.WriteByte(' ')
317 if len(ft.Results.List) == 1 && len(ft.Results.List[0].Names) == 0 {
318 buf.WriteString(renderNode(fset, ft.Results.List[0].Type))
319 } else {
320 buf.WriteByte('(')
321 for i, r := range ft.Results.List {
322 if i > 0 {
323 buf.WriteString(", ")
324 }
325 for j, name := range r.Names {
326 if j > 0 {
327 buf.WriteString(", ")
328 }
329 buf.WriteString(name.Name)
330 }
331 if r.Type != nil {
332 if len(r.Names) > 0 {
333 buf.WriteByte(' ')
334 }
335 buf.WriteString(renderNode(fset, r.Type))
336 }
337 }
338 buf.WriteByte(')')
339 }
340 }
341 return buf.String()
342 }
343
344 // emitFuncTypeIdents emits subtyped idents for function parameter and
345 // result names. Called from emitFuncDecl for each function/method.
346 func emitFuncTypeIdents(ch chan<- axiom.Element, ft *ast.FuncType) {
347 if ft == nil {
348 return
349 }
350 if ft.Params != nil {
351 for _, p := range ft.Params.List {
352 for _, name := range p.Names {
353 ch <- newHexElement("ident:param", name.Name)
354 }
355 }
356 }
357 if ft.Results != nil {
358 for _, r := range ft.Results.List {
359 for _, name := range r.Names {
360 ch <- newHexElement("ident:result", name.Name)
361 }
362 }
363 }
364 }
365
366 // declNameRe matches Go declaration keywords followed by a name.
367 var declNameRe = regexp.MustCompile(`(?m)^(?:var|type|func|const)\s+(\w+)`)
368
369 // emitDirectives scans raw source for //go: directives and emits them
370 // as "directive" elements. Each directive is associated with the next
371 // var/type/func/const declaration that follows it.
372 func emitDirectives(ch chan<- axiom.Element, src []byte) {
373 lines := strings.Split(string(src), "\n")
374 for i, line := range lines {
375 trimmed := strings.TrimSpace(line)
376 if !strings.HasPrefix(trimmed, "//go:") {
377 continue
378 }
379 // Find the associated declaration: scan forward for var/type/func/const.
380 assocName := ""
381 for j := i + 1; j < len(lines); j++ {
382 m := declNameRe.FindStringSubmatch(strings.TrimSpace(lines[j]))
383 if len(m) > 1 {
384 assocName = m[1]
385 break
386 }
387 }
388 ch <- newHexElement("directive", assocName+"\x00"+trimmed)
389 }
390 }
391