train.go raw
1 // Package train implements the training loop: pick a function with tests,
2 // hide the original, regenerate from description, run tests to verify
3 // functional equivalence.
4 //
5 // The organism proves it understands a function by reproducing it from its
6 // own English description. The tests are the oracle of truth — if they pass,
7 // the generated code is functionally equivalent.
8 //
9 // The Run function accepts a CodeGenerator interface so the generation
10 // strategy is pluggable (lattice-based, deterministic, etc.).
11 package train
12
13 import (
14 "context"
15 "fmt"
16 "os"
17 "os/exec"
18 "path/filepath"
19 "sort"
20 "strings"
21
22 "git.mleku.dev/mleku/dendrite/pkg/cartography"
23 )
24
25 // Target is a function selected for training.
26 type Target struct {
27 Entry *cartography.Entry
28 OrigSource string // original function source
29 Description string // English description from atlas
30 TestPkg string // package to test (e.g., "./describe/...")
31 TypeDefs string // type definitions from the same package that this function uses
32 TestSource string // actual test code for the function (so oracle can see expectations)
33 }
34
35 // Result captures the outcome of a training attempt.
36 type Result struct {
37 Target Target
38 Attempts int // how many oracle calls were needed
39 MaxAttempts int // limit
40 Passed bool // all tests passed
41 Generated string // the winning (or last) generated source
42 LastError string // last test failure output
43 }
44
45 // Config controls the training loop.
46 type Config struct {
47 MaxAttempts int // max oracle calls per target (default 5)
48 GoRoot string // path to Go installation
49 ProjectRoot string // path to project root
50 Verbose bool // print detailed output
51 }
52
53 // FindCandidates returns atlas entries that are good training targets:
54 // exported functions with tests, in packages that can be tested independently.
55 func FindCandidates(atlas *cartography.Atlas) []*cartography.Entry {
56 ids := sortedEntryIDs(atlas)
57 var candidates []*cartography.Entry
58 for _, id := range ids {
59 e := atlas.Entries[id]
60 if !e.HasTest {
61 continue
62 }
63 if !e.Exported {
64 continue
65 }
66 if e.Kind != "func" {
67 continue // start with standalone functions, not methods
68 }
69 if e.Description == "" {
70 continue
71 }
72 if e.Signature == "" {
73 continue
74 }
75 candidates = append(candidates, e)
76 }
77 return candidates
78 }
79
80 // CodeGenerator produces Go source from a query string.
81 // This replaces the former oracle dependency with a pluggable interface.
82 type CodeGenerator interface {
83 // GenerateCode takes a prompt and returns Go source code.
84 GenerateCode(ctx context.Context, query string) (string, error)
85 // SetGeneration advances the generation counter (for rate limiting).
86 SetGeneration(gen int)
87 }
88
89 // Run executes the training loop for a single target.
90 //
91 // Steps:
92 // 1. Copy project to temp dir
93 // 2. Remove the target function from the temp copy
94 // 3. Generate replacement from description via CodeGenerator
95 // 4. Write generated code to temp copy
96 // 5. Run tests — if pass, done
97 // 6. If fail, compose feedback query with test errors and retry
98 func Run(ctx context.Context, gen CodeGenerator, target Target, cfg Config) *Result {
99 if cfg.MaxAttempts <= 0 {
100 cfg.MaxAttempts = 5
101 }
102
103 result := &Result{
104 Target: target,
105 MaxAttempts: cfg.MaxAttempts,
106 }
107
108 // Create temp copy of project.
109 tmpDir, err := os.MkdirTemp("", "train-*")
110 if err != nil {
111 result.LastError = fmt.Sprintf("create temp dir: %v", err)
112 return result
113 }
114 defer os.RemoveAll(tmpDir)
115
116 if err := copyProject(cfg.ProjectRoot, tmpDir); err != nil {
117 result.LastError = fmt.Sprintf("copy project: %v", err)
118 return result
119 }
120
121 // Remove the target function from the temp copy and fix imports.
122 if err := removeFunction(tmpDir, target.Entry); err != nil {
123 result.LastError = fmt.Sprintf("remove function: %v", err)
124 return result
125 }
126
127 // Run goimports to clean up unused imports in the modified file.
128 goimportsCleanup(tmpDir, target.Entry.FilePath, cfg.GoRoot)
129
130 if cfg.Verbose {
131 fmt.Printf(" removed %s from temp copy\n", target.Entry.ID)
132 }
133
134 // Iterative generation loop.
135 var lastTestOutput string
136 for attempt := 1; attempt <= cfg.MaxAttempts; attempt++ {
137 result.Attempts = attempt
138
139 // Compose the generation query.
140 query := composeTrainQuery(target, lastTestOutput, attempt)
141
142 // Call generator.
143 gen.SetGeneration(attempt)
144 answer, err := gen.GenerateCode(ctx, query)
145 if err != nil {
146 result.LastError = fmt.Sprintf("generate call %d: %v", attempt, err)
147 if cfg.Verbose {
148 fmt.Printf(" attempt %d: generate error: %v\n", attempt, err)
149 }
150 continue
151 }
152
153 // Extract Go source from response.
154 // We use extractTrainSource instead of describe.ExtractGoSource
155 // because the target function may itself contain backtick markers
156 // that confuse the standard extraction regex.
157 goSource := extractTrainSource(answer)
158 if goSource == "" {
159 result.LastError = "no Go source in oracle response"
160 if cfg.Verbose {
161 fmt.Printf(" attempt %d: no source extracted\n", attempt)
162 }
163 continue
164 }
165
166 // Fix package declaration if oracle used wrong package.
167 goSource = fixPackage(goSource, target.Entry.Package)
168 result.Generated = goSource
169
170 if cfg.Verbose {
171 lines := strings.Split(goSource, "\n")
172 n := 10
173 if len(lines) < n {
174 n = len(lines)
175 }
176 fmt.Printf(" attempt %d: generated %d bytes (%d lines)\n", attempt, len(goSource), len(lines))
177 for _, line := range lines[:n] {
178 fmt.Printf(" > %s\n", line)
179 }
180 if len(lines) > n {
181 fmt.Printf(" > ... (%d more lines)\n", len(lines)-n)
182 }
183 }
184
185 // Write the generated code to the temp copy.
186 genPath := filepath.Join(tmpDir, target.Entry.FilePath)
187 // The generated code needs to go into the same package.
188 // Write as a new file alongside the original (with the function removed).
189 genFile := filepath.Join(filepath.Dir(genPath), "train_generated.go")
190 if err := os.WriteFile(genFile, []byte(goSource), 0o644); err != nil {
191 result.LastError = fmt.Sprintf("write generated: %v", err)
192 continue
193 }
194
195 // Run tests.
196 testPkg := target.TestPkg
197 if testPkg == "" {
198 testPkg = "./" + filepath.Dir(target.Entry.FilePath) + "/..."
199 }
200 testOutput, testErr := runTests(tmpDir, testPkg, cfg.GoRoot, target.Entry.TestFuncs)
201
202 if testErr == nil {
203 // Tests passed — functional equivalence achieved.
204 result.Passed = true
205 if cfg.Verbose {
206 fmt.Printf(" attempt %d: PASSED\n", attempt)
207 }
208 return result
209 }
210
211 lastTestOutput = testOutput
212 result.LastError = testOutput
213 if cfg.Verbose {
214 // Show first 5 lines of test output.
215 lines := strings.Split(testOutput, "\n")
216 n := 5
217 if len(lines) < n {
218 n = len(lines)
219 }
220 fmt.Printf(" attempt %d: FAILED\n", attempt)
221 for _, line := range lines[:n] {
222 fmt.Printf(" %s\n", line)
223 }
224 if len(lines) > n {
225 fmt.Printf(" ... (%d more lines)\n", len(lines)-n)
226 }
227 }
228
229 // Clean up generated file for next attempt.
230 os.Remove(genFile)
231 }
232
233 return result
234 }
235
236 // composeTrainQuery builds the oracle prompt for generating a function.
237 func composeTrainQuery(target Target, lastTestError string, attempt int) string {
238 var b strings.Builder
239
240 b.WriteString("Generate a Go function that is functionally equivalent to the following specification.\n\n")
241
242 // Include dependent type definitions so the oracle knows the data model.
243 if target.TypeDefs != "" {
244 b.WriteString("The following types are used by this function. They are already defined — do NOT redefine them.\n")
245 b.WriteString("Types from other packages are imported; use their package qualifier (e.g., axiom.Element).\n")
246 b.WriteString("IMPORTANT: Use the FULL import path shown in the comments (e.g., \"git.mleku.dev/mleku/dendrite/pkg/axiom\"), NOT a short path like \"axiom\".\n")
247 b.WriteString("Types from the same package are used directly (e.g., Description).\n\n")
248 b.WriteString(target.TypeDefs)
249 b.WriteString("\n")
250 }
251
252 // Description from atlas.
253 fmt.Fprintf(&b, "Function: %s\n", target.Entry.Name)
254 fmt.Fprintf(&b, "Package: %s\n", target.Entry.Package)
255 fmt.Fprintf(&b, "Signature: %s\n", target.Entry.Signature)
256 fmt.Fprintf(&b, "\nDescription: %s\n", target.Description)
257
258 if target.Entry.Contract != "" {
259 fmt.Fprintf(&b, "\nContract: %s\n", target.Entry.Contract)
260 }
261 if target.Entry.EdgeCases != "" {
262 fmt.Fprintf(&b, "\nEdge cases: %s\n", target.Entry.EdgeCases)
263 }
264 if target.Entry.DocComment != "" {
265 fmt.Fprintf(&b, "\nDoc comment: %s\n", target.Entry.DocComment)
266 }
267
268 // Test specifications — what the function must satisfy.
269 if len(target.Entry.TestFuncs) > 0 {
270 fmt.Fprintf(&b, "\nTests that must pass: %s\n", strings.Join(target.Entry.TestFuncs, ", "))
271 }
272 if len(target.Entry.TestPatterns) > 0 {
273 fmt.Fprintf(&b, "Test cases: %s\n", strings.Join(target.Entry.TestPatterns, ", "))
274 }
275 if target.Entry.Validation != "" {
276 fmt.Fprintf(&b, "Validation: %s\n", target.Entry.Validation)
277 }
278
279 // Include the actual test source if available, so the oracle can see
280 // the exact behavioral expectations.
281 if target.TestSource != "" {
282 b.WriteString("\nActual test code (for reference — your function must pass these):\n```go\n")
283 b.WriteString(target.TestSource)
284 b.WriteString("\n```\n")
285 }
286
287 // Parameter details.
288 if len(target.Entry.Params) > 0 {
289 b.WriteString("\nParameters:\n")
290 for _, p := range target.Entry.Params {
291 fmt.Fprintf(&b, " - %s %s", p.Name, p.Type)
292 if p.Semantic != "" {
293 fmt.Fprintf(&b, " (%s)", p.Semantic)
294 }
295 b.WriteString("\n")
296 }
297 }
298 if len(target.Entry.Returns) > 0 {
299 b.WriteString("Returns:\n")
300 for _, r := range target.Entry.Returns {
301 fmt.Fprintf(&b, " - %s", r.Type)
302 if r.IsError {
303 b.WriteString(" (error)")
304 }
305 if r.Semantic != "" {
306 fmt.Fprintf(&b, " — %s", r.Semantic)
307 }
308 b.WriteString("\n")
309 }
310 }
311
312 // Side effects.
313 if len(target.Entry.SideEffects) > 0 {
314 fmt.Fprintf(&b, "\nSide effects: %s\n", strings.Join(target.Entry.SideEffects, ", "))
315 }
316
317 // Dependencies — what other functions/types this calls.
318 if len(target.Entry.DependsOn) > 0 {
319 fmt.Fprintf(&b, "Dependencies: %s\n", strings.Join(target.Entry.DependsOn, ", "))
320 }
321
322 // Feedback from previous failed attempt.
323 if attempt > 1 && lastTestError != "" {
324 b.WriteString("\n--- PREVIOUS ATTEMPT FAILED ---\n")
325 b.WriteString("The following test failures occurred:\n\n")
326 // Only include FAIL lines and error messages — filter out PASS lines.
327 filtered := filterTestFailures(lastTestError)
328 b.WriteString(filtered)
329 b.WriteString("\n\nFix the issues and generate a corrected version.\n")
330 }
331
332 b.WriteString("\nRequirements:\n")
333 fmt.Fprintf(&b, "- Use package %s\n", target.Entry.Package)
334 b.WriteString("- Include all necessary imports\n")
335 b.WriteString("- The function must have the EXACT same signature\n")
336 b.WriteString("- Return the complete Go source inside a single ```go code block\n")
337 b.WriteString("- Include ONLY the function — no test code, no main function\n")
338 b.WriteString("- IMPORTANT: If the function body contains backtick characters, use raw string literals (backtick) carefully — the code block must still be valid\n")
339
340 return b.String()
341 }
342
343 // removeFunction removes a specific function from its source file in the
344 // temp project copy. It reads the file, finds the function by line number
345 // and brace matching, and writes the file back without it.
346 func removeFunction(tmpDir string, entry *cartography.Entry) error {
347 path := filepath.Join(tmpDir, entry.FilePath)
348 data, err := os.ReadFile(path)
349 if err != nil {
350 return err
351 }
352
353 lines := strings.Split(string(data), "\n")
354 start := entry.Line - 1 // 0-indexed
355 if start < 0 || start >= len(lines) {
356 return fmt.Errorf("line %d out of range (file has %d lines)", entry.Line, len(lines))
357 }
358
359 // Include the doc comment above the function.
360 for start > 0 && strings.HasPrefix(strings.TrimSpace(lines[start-1]), "//") {
361 start--
362 }
363
364 // Find the end by brace matching.
365 end := start
366 depth := 0
367 foundOpen := false
368 for i := start; i < len(lines); i++ {
369 for _, ch := range lines[i] {
370 if ch == '{' {
371 depth++
372 foundOpen = true
373 }
374 if ch == '}' {
375 depth--
376 if depth == 0 && foundOpen {
377 end = i + 1
378 goto found
379 }
380 }
381 }
382 }
383 // No braces found — single-line declaration or type.
384 end = start + 1
385
386 found:
387 // Remove lines [start:end].
388 var result []string
389 result = append(result, lines[:start]...)
390 result = append(result, lines[end:]...)
391
392 return os.WriteFile(path, []byte(strings.Join(result, "\n")), 0o644)
393 }
394
395 // runTests runs `go test` for a specific package in the temp directory.
396 // If testFuncs is non-empty, uses -run to filter to those specific tests.
397 // Returns the test output and any error.
398 func runTests(tmpDir, testPkg, goRoot string, testFuncs []string) (string, error) {
399 goBin := filepath.Join(goRoot, "bin", "go")
400 args := []string{"test", "-v", "-count=1"}
401 if len(testFuncs) > 0 {
402 args = append(args, "-run", "^("+strings.Join(testFuncs, "|")+")$")
403 }
404 args = append(args, testPkg)
405 cmd := exec.Command(goBin, args...)
406 cmd.Dir = tmpDir
407 cmd.Env = cleanGoEnv(goRoot)
408 out, err := cmd.CombinedOutput()
409 return string(out), err
410 }
411
412 // copyProject copies a Go project to a temp directory, skipping non-essential
413 // files for speed.
414 func copyProject(src, dst string) error {
415 return filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
416 if err != nil {
417 return err
418 }
419
420 rel, err := filepath.Rel(src, path)
421 if err != nil {
422 return err
423 }
424
425 if info.IsDir() {
426 base := filepath.Base(path)
427 switch base {
428 case ".git", "_output", "node_modules", "vendor",
429 ".svelte-kit", ".next", "dist", "build", "coverage":
430 return filepath.SkipDir
431 }
432 return os.MkdirAll(filepath.Join(dst, rel), info.Mode())
433 }
434
435 // Only copy Go source and module files.
436 ext := filepath.Ext(path)
437 if ext != ".go" && ext != ".mod" && ext != ".sum" {
438 return nil
439 }
440
441 data, err := os.ReadFile(path)
442 if err != nil {
443 return err
444 }
445 return os.WriteFile(filepath.Join(dst, rel), data, info.Mode())
446 })
447 }
448
449 // cleanGoEnv returns an environment with GOROOT, PATH, and GOTOOLCHAIN set.
450 func cleanGoEnv(root string) []string {
451 env := os.Environ()
452 clean := make([]string, 0, len(env)+3)
453 for _, e := range env {
454 if strings.HasPrefix(e, "GOROOT=") ||
455 strings.HasPrefix(e, "GOTOOLCHAIN=") ||
456 strings.HasPrefix(e, "PATH=") {
457 continue
458 }
459 clean = append(clean, e)
460 }
461 clean = append(clean,
462 "GOROOT="+root,
463 "GOTOOLCHAIN=local",
464 "PATH="+filepath.Join(root, "bin")+":"+os.Getenv("PATH"),
465 )
466 return clean
467 }
468
469 // goimportsCleanup removes unused imports from a Go file after a function
470 // has been deleted. It does a simple text scan: finds all import paths,
471 // checks if each is referenced in the remaining code, and removes unreferenced ones.
472 func goimportsCleanup(tmpDir, relPath, goRoot string) {
473 absPath := filepath.Join(tmpDir, relPath)
474 data, err := os.ReadFile(absPath)
475 if err != nil {
476 return
477 }
478
479 source := string(data)
480 lines := strings.Split(source, "\n")
481
482 // Find import block boundaries and individual imports.
483 type importLine struct {
484 lineIdx int
485 path string // e.g., "go/ast"
486 alias string // e.g., "ast" or custom alias
487 }
488
489 var imports []importLine
490 inImportBlock := false
491 importBlockStart := -1
492 importBlockEnd := -1
493
494 for i, line := range lines {
495 trimmed := strings.TrimSpace(line)
496
497 if trimmed == "import (" {
498 inImportBlock = true
499 importBlockStart = i
500 continue
501 }
502 if inImportBlock && trimmed == ")" {
503 inImportBlock = false
504 importBlockEnd = i
505 continue
506 }
507 if inImportBlock && trimmed != "" && !strings.HasPrefix(trimmed, "//") {
508 // Parse import line: could be `"path"` or `alias "path"`
509 path := ""
510 alias := ""
511 parts := strings.Fields(trimmed)
512 for _, p := range parts {
513 clean := strings.Trim(p, `"`)
514 if clean != p { // it was quoted — this is the path
515 path = clean
516 } else {
517 alias = p
518 }
519 }
520 if path != "" {
521 if alias == "" {
522 alias = importAlias(path)
523 }
524 imports = append(imports, importLine{lineIdx: i, path: path, alias: alias})
525 }
526 }
527 }
528
529 if len(imports) == 0 || importBlockStart < 0 {
530 return
531 }
532
533 // Build the code text WITHOUT import block for searching.
534 var codeLines []string
535 for i, line := range lines {
536 if i >= importBlockStart && i <= importBlockEnd {
537 continue
538 }
539 codeLines = append(codeLines, line)
540 }
541 codeText := strings.Join(codeLines, "\n")
542
543 // Check which imports are used by looking for the alias as a package
544 // qualifier. We need word-boundary awareness: "ast." must not match
545 // "fast." or "last.". A package qualifier in Go is always preceded by
546 // a non-alphanumeric character (space, tab, paren, star, ampersand, etc.)
547 // or starts a line.
548 var unusedLineIdxs []int
549 for _, imp := range imports {
550 if imp.alias == "_" {
551 continue
552 }
553 used := isImportUsed(codeText, imp.alias)
554 if !used {
555 unusedLineIdxs = append(unusedLineIdxs, imp.lineIdx)
556 }
557 }
558
559 if len(unusedLineIdxs) == 0 {
560 return
561 }
562
563 // Remove unused import lines.
564 skip := make(map[int]bool)
565 for _, idx := range unusedLineIdxs {
566 skip[idx] = true
567 }
568
569 var result []string
570 for i, line := range lines {
571 if skip[i] {
572 continue
573 }
574 result = append(result, line)
575 }
576
577 os.WriteFile(absPath, []byte(strings.Join(result, "\n")), 0o644)
578 }
579
580 // filterTestFailures extracts only the failure-relevant lines from test output.
581 // Removes PASS lines and keeps FAIL lines, error messages, and compiler errors.
582 func filterTestFailures(output string) string {
583 lines := strings.Split(output, "\n")
584 var result []string
585 inFailBlock := false
586
587 for _, line := range lines {
588 trimmed := strings.TrimSpace(line)
589
590 // Always include compiler errors.
591 if strings.HasPrefix(trimmed, "#") || strings.Contains(line, ": undefined:") ||
592 strings.Contains(line, "imported and not used") ||
593 strings.Contains(line, "redeclared") {
594 result = append(result, line)
595 continue
596 }
597
598 // Skip PASS lines.
599 if strings.HasPrefix(trimmed, "--- PASS:") || strings.HasPrefix(trimmed, "=== PAUSE") {
600 inFailBlock = false
601 continue
602 }
603
604 // Include FAIL lines and their context.
605 if strings.HasPrefix(trimmed, "--- FAIL:") || strings.HasPrefix(trimmed, "FAIL") {
606 result = append(result, line)
607 inFailBlock = true
608 continue
609 }
610
611 // Include test error output (indented lines after a RUN or FAIL).
612 if strings.HasPrefix(trimmed, "=== RUN") {
613 // Check if this test fails — peek ahead is hard, so include the RUN.
614 inFailBlock = true
615 result = append(result, line)
616 continue
617 }
618
619 // Include assertion failures and error messages.
620 if inFailBlock && trimmed != "" {
621 result = append(result, line)
622 }
623
624 // Cap output.
625 if len(result) >= 40 {
626 result = append(result, "... (truncated)")
627 break
628 }
629 }
630
631 return strings.Join(result, "\n")
632 }
633
634 // isImportUsed checks if a package alias is used as a qualifier in Go code.
635 // Handles word boundaries: "ast." must not match "fast." or "last.".
636 func isImportUsed(code, alias string) bool {
637 needle := alias + "."
638 idx := 0
639 for {
640 pos := strings.Index(code[idx:], needle)
641 if pos < 0 {
642 return false
643 }
644 absPos := idx + pos
645 // Check that the character before the match is not alphanumeric.
646 if absPos == 0 {
647 return true // starts at beginning of code
648 }
649 prev := code[absPos-1]
650 if !isAlphaNum(prev) {
651 return true // word boundary before alias
652 }
653 idx = absPos + len(needle)
654 if idx >= len(code) {
655 return false
656 }
657 }
658 }
659
660 func isAlphaNum(b byte) bool {
661 return (b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z') || (b >= '0' && b <= '9') || b == '_'
662 }
663
664 // importAlias derives the Go package name from an import path.
665 // Handles versioned module paths: "github.com/foo/bar/v2" → "bar",
666 // "github.com/btcsuite/btcd/btcec/v2" → "btcec".
667 func importAlias(path string) string {
668 parts := strings.Split(path, "/")
669 last := parts[len(parts)-1]
670 // If the last component is a Go module version suffix (v2, v3, etc.),
671 // use the second-to-last component.
672 if len(parts) >= 2 && len(last) >= 2 && last[0] == 'v' && last[1] >= '0' && last[1] <= '9' {
673 return parts[len(parts)-2]
674 }
675 return last
676 }
677
678 // extractTrainSource extracts Go source from an oracle response, handling
679 // the case where the generated function body itself contains backtick markers
680 // (e.g., when regenerating ExtractGoSource). Unlike describe.ExtractGoSource
681 // which uses non-greedy regex, this finds code blocks by locating the first
682 // ```go marker and the LAST ``` marker, which handles nested backticks.
683 func extractTrainSource(answer string) string {
684 // Find the opening marker.
685 openIdx := strings.Index(answer, "```go")
686 if openIdx < 0 {
687 openIdx = strings.Index(answer, "```")
688 if openIdx < 0 {
689 // No code fences — if it starts with "package", use the whole thing.
690 trimmed := strings.TrimSpace(answer)
691 if strings.HasPrefix(trimmed, "package ") {
692 return trimmed
693 }
694 return ""
695 }
696 }
697
698 // Move past the opening marker line.
699 startOfCode := strings.Index(answer[openIdx:], "\n")
700 if startOfCode < 0 {
701 return ""
702 }
703 startOfCode += openIdx + 1
704
705 // Find the LAST ``` in the response (closing marker).
706 // This handles cases where the generated code contains backticks.
707 lastClose := strings.LastIndex(answer, "```")
708 if lastClose <= openIdx {
709 // No closing marker — use everything after opening.
710 return strings.TrimSpace(answer[startOfCode:])
711 }
712
713 source := answer[startOfCode:lastClose]
714 source = strings.TrimSpace(source)
715
716 // Safety: strip any residual opening backtick markers at the start.
717 for strings.HasPrefix(source, "```") {
718 nl := strings.Index(source, "\n")
719 if nl < 0 {
720 break
721 }
722 source = strings.TrimSpace(source[nl+1:])
723 }
724
725 return source
726 }
727
728 // fixPackage ensures the generated source uses the correct package name.
729 // If the oracle produced "package main" but we need "package describe",
730 // this fixes it.
731 func fixPackage(source, wantPkg string) string {
732 lines := strings.Split(source, "\n")
733 for i, line := range lines {
734 trimmed := strings.TrimSpace(line)
735 if strings.HasPrefix(trimmed, "package ") {
736 parts := strings.Fields(trimmed)
737 if len(parts) >= 2 && parts[1] != wantPkg {
738 lines[i] = "package " + wantPkg
739 }
740 break
741 }
742 }
743 return strings.Join(lines, "\n")
744 }
745
746 // ExtractTypeDefs finds type definitions that the target function's parameters
747 // or returns reference. Handles both same-package types (e.g., Description)
748 // and cross-package types (e.g., axiom.Element, lattice.Lattice).
749 // Also extracts method signatures for cross-package types so the oracle knows
750 // the API surface (e.g., Lattice.Nodes(), Node.Bond()).
751 // Reads actual source from disk to get full struct/interface definitions.
752 func ExtractTypeDefs(entry *cartography.Entry, atlas *cartography.Atlas, projectRoot string) string {
753 var b strings.Builder
754
755 // Read module path from go.mod for constructing full import paths.
756 modulePath := readModulePath(projectRoot)
757
758 // Collect type references from params and returns.
759 // samePackage: {"Description": true}
760 // crossPackage: {"axiom.Element": true, "lattice.Lattice": true}
761 samePackage := make(map[string]bool)
762 crossPackage := make(map[string]bool)
763
764 for _, p := range entry.Params {
765 classifyType(p.Type, entry.Package, samePackage, crossPackage)
766 }
767 for _, r := range entry.Returns {
768 classifyType(r.Type, entry.Package, samePackage, crossPackage)
769 }
770
771 // Iteratively expand types: for each type found, also collect types
772 // referenced by its struct fields and method return types.
773 // Run up to 3 rounds to follow chains like Lattice → Node → axiom.Constraint.
774 for range 3 {
775 before := len(samePackage) + len(crossPackage)
776 expandFieldTypes(samePackage, crossPackage, entry.Package, atlas)
777 after := len(samePackage) + len(crossPackage)
778 if after == before {
779 break // no new types discovered
780 }
781 }
782
783 if len(samePackage) == 0 && len(crossPackage) == 0 {
784 return ""
785 }
786
787 // Collect already-emitted type IDs to avoid duplicates.
788 emitted := make(map[string]bool)
789
790 // Same-package types: emit ALL exported type definitions from the package.
791 // The oracle needs to see every type it might construct or reference.
792 // Same-package types are always accessible, so this is safe.
793 ids := sortedEntryIDs(atlas)
794 for _, id := range ids {
795 e := atlas.Entries[id]
796 if e.Package != entry.Package {
797 continue
798 }
799 if e.Kind != "type" && e.Kind != "interface" {
800 continue
801 }
802 if !e.Exported {
803 continue
804 }
805 if emitted[e.ID] {
806 continue
807 }
808 src := readTypeSource(e, projectRoot)
809 if src != "" {
810 fmt.Fprintf(&b, "// (same package %s)\n%s\n\n", e.Package, src)
811 emitted[e.ID] = true
812 }
813 }
814
815 // Cross-package types: emit struct/interface definitions + method signatures.
816 crossNames := make([]string, 0, len(crossPackage))
817 for qn := range crossPackage {
818 crossNames = append(crossNames, qn)
819 }
820 sort.Strings(crossNames)
821 for _, qualifiedName := range crossNames {
822 parts := strings.SplitN(qualifiedName, ".", 2)
823 if len(parts) != 2 {
824 continue
825 }
826 pkg, name := parts[0], parts[1]
827
828 // Find the atlas entry for this type.
829 for _, id := range ids {
830 e := atlas.Entries[id]
831 if e.Package != pkg || e.Name != name {
832 continue
833 }
834 if e.Kind != "type" && e.Kind != "interface" {
835 continue
836 }
837 if emitted[e.ID] {
838 continue
839 }
840 src := readTypeSource(e, projectRoot)
841 fullImport := modulePath + "/" + filepath.Dir(e.FilePath)
842 if src != "" {
843 fmt.Fprintf(&b, "// (from package %s — import path: %q, do NOT redefine)\n%s\n", pkg, fullImport, src)
844 emitted[e.ID] = true
845 } else if e.Description != "" {
846 fmt.Fprintf(&b, "// (from package %s — import path: %q) %s\n", pkg, fullImport, e.Description)
847 emitted[e.ID] = true
848 }
849
850 // Append exported method signatures for this type.
851 methods := collectMethods(pkg, name, atlas)
852 if len(methods) > 0 {
853 fmt.Fprintf(&b, "// Exported methods on %s.%s:\n", pkg, name)
854 for _, sig := range methods {
855 fmt.Fprintf(&b, "// %s\n", sig)
856 }
857 }
858 b.WriteString("\n")
859 break
860 }
861 }
862
863 return b.String()
864 }
865
866 // expandFieldTypes looks at each already-collected type's struct fields AND
867 // method return types in the atlas and adds any new type references.
868 // One level only — prevents explosion. This ensures that if Lattice has
869 // method Nodes() []*Node, the Node type is also collected.
870 func expandFieldTypes(samePackage, crossPackage map[string]bool, currentPkg string, atlas *cartography.Atlas) {
871 // Snapshot current sets so we don't iterate while modifying.
872 sameNames := make([]string, 0, len(samePackage))
873 for k := range samePackage {
874 sameNames = append(sameNames, k)
875 }
876 sort.Strings(sameNames)
877
878 crossNames := make([]string, 0, len(crossPackage))
879 for k := range crossPackage {
880 crossNames = append(crossNames, k)
881 }
882 sort.Strings(crossNames)
883
884 ids := sortedEntryIDs(atlas)
885
886 // For each same-package type, collect types from struct fields.
887 for _, name := range sameNames {
888 for _, id := range ids {
889 e := atlas.Entries[id]
890 if e.Package != currentPkg || e.Name != name {
891 continue
892 }
893 if e.Kind != "type" {
894 continue
895 }
896 for _, f := range e.Params {
897 classifyType(f.Type, currentPkg, samePackage, crossPackage)
898 }
899 break
900 }
901 }
902
903 // For each cross-package type, collect types from struct fields
904 // AND from method return types.
905 for _, qualifiedName := range crossNames {
906 parts := strings.SplitN(qualifiedName, ".", 2)
907 if len(parts) != 2 {
908 continue
909 }
910 pkg, name := parts[0], parts[1]
911
912 // Struct fields.
913 for _, id := range ids {
914 e := atlas.Entries[id]
915 if e.Package != pkg || e.Name != name || e.Kind != "type" {
916 continue
917 }
918 for _, f := range e.Params {
919 classifyType(f.Type, currentPkg, samePackage, crossPackage)
920 }
921 break
922 }
923
924 // Method return types: if Lattice.Nodes() returns []*Node,
925 // add lattice.Node to crossPackage so its methods are also extracted.
926 // Note: return types are recorded unqualified within their own package,
927 // so "[]*Node" in a lattice method means "lattice.Node" for our purposes.
928 for _, id := range ids {
929 e := atlas.Entries[id]
930 if e.Package != pkg || e.Kind != "method" || !e.Exported {
931 continue
932 }
933 recv := strings.TrimPrefix(e.Receiver, "*")
934 if recv != name {
935 continue
936 }
937 for _, r := range e.Returns {
938 // Classify relative to the method's own package so that
939 // unqualified types like "Node" resolve to same-package
940 // of the method (i.e., "lattice"), then we promote them
941 // to cross-package from the target's perspective.
942 methodSame := make(map[string]bool)
943 methodCross := make(map[string]bool)
944 classifyType(r.Type, pkg, methodSame, methodCross)
945 // Same-package of the method = cross-package of the target.
946 for typeName := range methodSame {
947 if pkg != currentPkg {
948 crossPackage[pkg+"."+typeName] = true
949 } else {
950 samePackage[typeName] = true
951 }
952 }
953 for k := range methodCross {
954 crossPackage[k] = true
955 }
956 }
957 }
958 }
959 }
960
961 // collectMethods finds all exported method signatures for a given type
962 // in the atlas. Returns signature strings like "Nodes() []*Node".
963 func collectMethods(pkg, typeName string, atlas *cartography.Atlas) []string {
964 ids := sortedEntryIDs(atlas)
965 var sigs []string
966 // Match methods whose receiver is the type (pointer or value).
967 for _, id := range ids {
968 e := atlas.Entries[id]
969 if e.Package != pkg || e.Kind != "method" {
970 continue
971 }
972 if !e.Exported {
973 continue
974 }
975 // Receiver is stored as "*Lattice" or "Lattice".
976 recv := strings.TrimPrefix(e.Receiver, "*")
977 if recv != typeName {
978 continue
979 }
980 // Use the signature, stripping the "func " prefix and receiver.
981 // The atlas Signature looks like: "func (l *Lattice) Nodes() []*Node"
982 // We want just: "Nodes() []*Node"
983 sig := e.Signature
984 if i := strings.Index(sig, ") "); i >= 0 {
985 sig = strings.TrimSpace(sig[i+2:])
986 }
987 sigs = append(sigs, sig)
988 }
989 return sigs
990 }
991
992 // sortedEntryIDs returns atlas entry IDs in sorted order for deterministic iteration.
993 func sortedEntryIDs(atlas *cartography.Atlas) []string {
994 ids := make([]string, 0, len(atlas.Entries))
995 for id := range atlas.Entries {
996 ids = append(ids, id)
997 }
998 sort.Strings(ids)
999 return ids
1000 }
1001
1002 // classifyType parses a Go type expression and classifies referenced custom
1003 // types as same-package or cross-package.
1004 // "*RevenueSummary" → samePackage["RevenueSummary"]
1005 // "[]axiom.Element" → crossPackage["axiom.Element"]
1006 // "*lattice.Lattice" → crossPackage["lattice.Lattice"]
1007 func classifyType(typ, currentPkg string, samePackage, crossPackage map[string]bool) {
1008 // Strip pointer, slice, variadic, map prefixes.
1009 t := typ
1010 t = strings.TrimPrefix(t, "*")
1011 t = strings.TrimPrefix(t, "[]")
1012 t = strings.TrimPrefix(t, "...")
1013 t = strings.TrimPrefix(t, "*") // **T
1014 // Handle map types: map[K]V — extract V.
1015 if strings.HasPrefix(t, "map[") {
1016 if i := strings.LastIndex(t, "]"); i >= 0 {
1017 t = t[i+1:]
1018 t = strings.TrimPrefix(t, "*")
1019 }
1020 }
1021
1022 if t == "" || isBuiltinType(t) {
1023 return
1024 }
1025
1026 // Check for package qualifier: "axiom.Element"
1027 if i := strings.Index(t, "."); i > 0 {
1028 crossPackage[t] = true
1029 } else if t[0] >= 'A' && t[0] <= 'Z' {
1030 // Unqualified uppercase → same package type.
1031 samePackage[t] = true
1032 }
1033 }
1034
1035 // isBuiltinType returns true for Go builtin types.
1036 func isBuiltinType(name string) bool {
1037 switch name {
1038 case "string", "int", "int8", "int16", "int32", "int64",
1039 "uint", "uint8", "uint16", "uint32", "uint64",
1040 "float32", "float64", "bool", "byte", "rune",
1041 "error", "any", "interface{}", "uintptr",
1042 "context.Context": // treat as builtin — everyone knows it
1043 return true
1044 }
1045 return false
1046 }
1047
1048 // readTypeSource reads the actual type definition from source.
1049 func readTypeSource(entry *cartography.Entry, projectRoot string) string {
1050 path := filepath.Join(projectRoot, entry.FilePath)
1051 data, err := os.ReadFile(path)
1052 if err != nil {
1053 return ""
1054 }
1055
1056 lines := strings.Split(string(data), "\n")
1057 start := entry.Line - 1
1058 if start < 0 || start >= len(lines) {
1059 return ""
1060 }
1061
1062 // Find the end by brace matching.
1063 depth := 0
1064 for i := start; i < len(lines) && i < start+50; i++ {
1065 for _, ch := range lines[i] {
1066 if ch == '{' {
1067 depth++
1068 }
1069 if ch == '}' {
1070 depth--
1071 if depth == 0 {
1072 return strings.Join(lines[start:i+1], "\n")
1073 }
1074 }
1075 }
1076 }
1077
1078 // No braces — single-line type.
1079 return lines[start]
1080 }
1081
1082 // readModulePath reads the Go module path from go.mod.
1083 func readModulePath(projectRoot string) string {
1084 data, err := os.ReadFile(filepath.Join(projectRoot, "go.mod"))
1085 if err != nil {
1086 return ""
1087 }
1088 for _, line := range strings.Split(string(data), "\n") {
1089 line = strings.TrimSpace(line)
1090 if strings.HasPrefix(line, "module ") {
1091 return strings.TrimSpace(strings.TrimPrefix(line, "module"))
1092 }
1093 }
1094 return ""
1095 }
1096
1097 // firstSentence returns the first sentence of a string.
1098 func firstSentence(s string) string {
1099 s = strings.TrimSpace(s)
1100 if i := strings.IndexByte(s, '.'); i >= 0 && i < 120 {
1101 return s[:i+1]
1102 }
1103 if len(s) > 120 {
1104 return s[:120] + "..."
1105 }
1106 return s
1107 }
1108
1109 // ExtractTestSource reads the test functions for an entry from the test file.
1110 func ExtractTestSource(entry *cartography.Entry, projectRoot string) string {
1111 if entry.TestFile == "" || len(entry.TestFuncs) == 0 {
1112 return ""
1113 }
1114
1115 path := filepath.Join(projectRoot, entry.TestFile)
1116 data, err := os.ReadFile(path)
1117 if err != nil {
1118 return ""
1119 }
1120
1121 source := string(data)
1122 lines := strings.Split(source, "\n")
1123 var result []string
1124
1125 for _, testFunc := range entry.TestFuncs {
1126 // Find the test function in the file.
1127 prefix := "func " + testFunc + "("
1128 for i, line := range lines {
1129 if !strings.Contains(line, prefix) {
1130 continue
1131 }
1132 // Found start — extract until closing brace.
1133 depth := 0
1134 for j := i; j < len(lines) && j < i+100; j++ {
1135 result = append(result, lines[j])
1136 for _, ch := range lines[j] {
1137 if ch == '{' {
1138 depth++
1139 }
1140 if ch == '}' {
1141 depth--
1142 if depth == 0 {
1143 goto nextFunc
1144 }
1145 }
1146 }
1147 }
1148 nextFunc:
1149 result = append(result, "")
1150 break
1151 }
1152 }
1153
1154 return strings.Join(result, "\n")
1155 }
1156
1157 // BuildDescription composes a rich English description of a function
1158 // from its atlas entry, suitable for training. This is what the oracle
1159 // receives instead of the source code.
1160 func BuildDescription(entry *cartography.Entry) string {
1161 var b strings.Builder
1162
1163 // Start with the atlas description.
1164 if entry.Description != "" {
1165 b.WriteString(entry.Description)
1166 }
1167
1168 // Add contract details.
1169 if entry.Contract != "" {
1170 fmt.Fprintf(&b, "\n\nContract: %s", entry.Contract)
1171 }
1172
1173 // Add edge cases.
1174 if entry.EdgeCases != "" {
1175 fmt.Fprintf(&b, "\n\nEdge cases: %s", entry.EdgeCases)
1176 }
1177
1178 // Add doc comment (original author's intent).
1179 if entry.DocComment != "" {
1180 fmt.Fprintf(&b, "\n\nDoc: %s", entry.DocComment)
1181 }
1182
1183 return b.String()
1184 }
1185