ast.go raw
1 package main
2
3 import (
4 "fmt"
5 "go/ast"
6 "go/token"
7 "os"
8 "path/filepath"
9 "strings"
10 )
11
12 type parsedFile struct {
13 name string
14 origSrc []byte // original .mx source before text rewrites
15 src []byte // after text rewrites (parseable Go)
16 file *ast.File
17 }
18
19 func extractAST(fset *token.FileSet, pf parsedFile, fileIndex int, outdir string) (FileManifest, error) {
20 if err := os.MkdirAll(outdir, 0755); err != nil {
21 return FileManifest{}, err
22 }
23
24 // Copy verbatim original source.
25 if err := os.WriteFile(filepath.Join(outdir, pf.name), pf.origSrc, 0644); err != nil {
26 return FileManifest{}, err
27 }
28
29 segDir := filepath.Join(outdir, pf.name+".segments")
30 if err := os.MkdirAll(segDir, 0755); err != nil {
31 return FileManifest{}, err
32 }
33
34 astDir := filepath.Join(outdir, "ast")
35 if err := os.MkdirAll(astDir, 0755); err != nil {
36 return FileManifest{}, err
37 }
38
39 fm := FileManifest{
40 Path: pf.name,
41 FileIndex: fileIndex,
42 }
43
44 // Use the rewritten source for segment extraction (AST positions match rewritten bytes).
45 segments := segmentFile(fset, pf.file, pf.src, fileIndex)
46 for si, seg := range segments {
47 segFilename := seg.Filename()
48 segPath := filepath.Join(segDir, segFilename)
49 if err := os.WriteFile(segPath, seg.Data, 0644); err != nil {
50 return FileManifest{}, fmt.Errorf("write segment %s: %w", segPath, err)
51 }
52
53 // Generate AST dump for declarations (si=0 is pkg clause, decls start at si=1).
54 astDumpFile := ""
55 declIdx := si - 1
56 if declIdx >= 0 && declIdx < len(pf.file.Decls) {
57 dump := dumpAST(fset, pf.file.Decls[declIdx])
58 if len(dump) > 0 {
59 astDumpFile = strings.TrimSuffix(segFilename, ".mx") + ".ast"
60 astPath := filepath.Join(astDir, astDumpFile)
61 if err := os.WriteFile(astPath, []byte(dump), 0644); err != nil {
62 return FileManifest{}, fmt.Errorf("write ast dump %s: %w", astPath, err)
63 }
64 }
65 }
66
67 sm := SegmentManifest{
68 ID: strings.TrimSuffix(segFilename, ".mx"),
69 Kind: seg.Kind,
70 Name: seg.Name,
71 ASTFile: segFilename,
72 ASTDump: astDumpFile,
73 SizeBytes: len(seg.Data),
74 SourcePos: SourcePos{Line: seg.StartLine, EndLine: seg.EndLine},
75 Comments: seg.Comments,
76 Generated: false,
77 HasSpawn: seg.HasSpawn,
78 }
79 fm.Segments = append(fm.Segments, sm)
80 }
81
82 return fm, nil
83 }
84
85 type segment struct {
86 FileIndex int
87 Seq int
88 Line int
89 Kind string
90 Name string
91 Data []byte
92 StartLine int
93 EndLine int
94 Comments []string
95 HasSpawn bool
96 }
97
98 func (s *segment) Filename() string {
99 return fmt.Sprintf("f%02d_%03d_L%d_%s_%s.mx", s.FileIndex, s.Seq, s.Line, s.Kind, sanitizeName(s.Name))
100 }
101
102 func sanitizeName(name string) string {
103 name = strings.ReplaceAll(name, "/", "_")
104 name = strings.ReplaceAll(name, " ", "_")
105 name = strings.ReplaceAll(name, "*", "_")
106 name = strings.ReplaceAll(name, "(", "")
107 name = strings.ReplaceAll(name, ")", "")
108 return name
109 }
110
111 func segmentFile(fset *token.FileSet, file *ast.File, src []byte, fileIndex int) []segment {
112 tokFile := fset.File(file.Pos())
113 if tokFile == nil {
114 return nil
115 }
116
117 var segments []segment
118 seq := 0
119
120 prevEnd := 0
121
122 // Package clause segment.
123 {
124 seq++
125 pkgEnd := tokFile.Offset(file.Name.End())
126 segEnd := skipToEndOfLine(src, pkgEnd)
127
128 var comments []string
129 if file.Doc != nil {
130 for _, c := range file.Doc.List {
131 comments = append(comments, c.Text)
132 }
133 }
134
135 startLine := 1
136 endLine := fset.Position(file.Name.End()).Line
137
138 segments = append(segments, segment{
139 FileIndex: fileIndex,
140 Seq: seq,
141 Line: startLine,
142 Kind: "pkg",
143 Name: file.Name.Name,
144 Data: src[0:segEnd],
145 StartLine: startLine,
146 EndLine: endLine,
147 Comments: comments,
148 })
149 prevEnd = segEnd
150 }
151
152 for i, decl := range file.Decls {
153 seq++
154
155 declStart := tokFile.Offset(decl.Pos())
156 declEnd := tokFile.Offset(decl.End())
157
158 docStart := declStart
159 switch d := decl.(type) {
160 case *ast.FuncDecl:
161 if d.Doc != nil {
162 ds := tokFile.Offset(d.Doc.Pos())
163 if ds < docStart {
164 docStart = ds
165 }
166 }
167 case *ast.GenDecl:
168 if d.Doc != nil {
169 ds := tokFile.Offset(d.Doc.Pos())
170 if ds < docStart {
171 docStart = ds
172 }
173 }
174 }
175
176 segStart := prevEnd
177 if docStart < segStart {
178 segStart = docStart
179 }
180
181 segEnd := declEnd
182 if i+1 < len(file.Decls) {
183 nextStart := offsetOfDecl(tokFile, file.Decls[i+1])
184 segEnd = nextStart
185 } else {
186 segEnd = len(src)
187 }
188
189 kind, name := classifyDecl(decl)
190 startLine := fset.Position(decl.Pos()).Line
191 endLine := fset.Position(decl.End()).Line
192
193 var comments []string
194 switch d := decl.(type) {
195 case *ast.FuncDecl:
196 if d.Doc != nil {
197 for _, c := range d.Doc.List {
198 comments = append(comments, c.Text)
199 }
200 }
201 case *ast.GenDecl:
202 if d.Doc != nil {
203 for _, c := range d.Doc.List {
204 comments = append(comments, c.Text)
205 }
206 }
207 }
208
209 hasSpawn := false
210 if fd, ok := decl.(*ast.FuncDecl); ok && fd.Body != nil {
211 hasSpawn = detectSpawn(fd.Body)
212 }
213
214 segments = append(segments, segment{
215 FileIndex: fileIndex,
216 Seq: seq,
217 Line: startLine,
218 Kind: kind,
219 Name: name,
220 Data: src[segStart:segEnd],
221 StartLine: startLine,
222 EndLine: endLine,
223 Comments: comments,
224 HasSpawn: hasSpawn,
225 })
226 prevEnd = segEnd
227 }
228
229 return segments
230 }
231
232 func offsetOfDecl(tokFile *token.File, decl ast.Decl) int {
233 start := tokFile.Offset(decl.Pos())
234 switch d := decl.(type) {
235 case *ast.FuncDecl:
236 if d.Doc != nil {
237 ds := tokFile.Offset(d.Doc.Pos())
238 if ds < start {
239 return ds
240 }
241 }
242 case *ast.GenDecl:
243 if d.Doc != nil {
244 ds := tokFile.Offset(d.Doc.Pos())
245 if ds < start {
246 return ds
247 }
248 }
249 }
250 return start
251 }
252
253 func classifyDecl(decl ast.Decl) (kind, name string) {
254 switch d := decl.(type) {
255 case *ast.FuncDecl:
256 if d.Recv != nil && len(d.Recv.List) > 0 {
257 recv := exprName(d.Recv.List[0].Type)
258 return "method", "(" + recv + ")." + d.Name.Name
259 }
260 return "func", d.Name.Name
261 case *ast.GenDecl:
262 switch d.Tok {
263 case token.IMPORT:
264 return "import", "import"
265 case token.CONST:
266 if len(d.Specs) > 0 {
267 if vs, ok := d.Specs[0].(*ast.ValueSpec); ok && len(vs.Names) > 0 {
268 return "const", vs.Names[0].Name
269 }
270 }
271 return "const", "const"
272 case token.TYPE:
273 if len(d.Specs) > 0 {
274 if ts, ok := d.Specs[0].(*ast.TypeSpec); ok {
275 return "type", ts.Name.Name
276 }
277 }
278 return "type", "type"
279 case token.VAR:
280 if len(d.Specs) > 0 {
281 if vs, ok := d.Specs[0].(*ast.ValueSpec); ok && len(vs.Names) > 0 {
282 return "var", vs.Names[0].Name
283 }
284 }
285 return "var", "var"
286 }
287 }
288 return "unknown", "unknown"
289 }
290
291 func exprName(e ast.Expr) string {
292 switch t := e.(type) {
293 case *ast.Ident:
294 return t.Name
295 case *ast.StarExpr:
296 return "*" + exprName(t.X)
297 case *ast.IndexExpr:
298 return exprName(t.X)
299 case *ast.IndexListExpr:
300 return exprName(t.X)
301 }
302 return "?"
303 }
304
305 func detectSpawn(body *ast.BlockStmt) bool {
306 found := false
307 ast.Inspect(body, func(n ast.Node) bool {
308 if found {
309 return false
310 }
311 call, ok := n.(*ast.CallExpr)
312 if !ok {
313 return true
314 }
315 if ident, ok := call.Fun.(*ast.Ident); ok && ident.Name == "spawn" {
316 found = true
317 return false
318 }
319 return true
320 })
321 return found
322 }
323
324 func skipToEndOfLine(src []byte, off int) int {
325 for off < len(src) {
326 if src[off] == '\n' {
327 return off + 1
328 }
329 off++
330 }
331 return off
332 }
333