fitness.go raw

   1  // Package fitness evaluates how closely emitted code reproduces the original.
   2  //
   3  // Three dimensions:
   4  //   - Source: structural similarity of Go AST (types, functions, methods)
   5  //   - Binary: size ratio and shared content of compiled binaries
   6  //   - Behavior: same inputs → same outputs
   7  //
   8  // The fitness score feeds back into the spore, driving evolution:
   9  // generations with higher fitness produce spores that nucleate
  10  // better-adapted lattices.
  11  package fitness
  12  
  13  import (
  14  	"bytes"
  15  	"go/ast"
  16  	"go/parser"
  17  	"go/token"
  18  	"os"
  19  	"os/exec"
  20  	"path/filepath"
  21  	"strings"
  22  
  23  	"git.mleku.dev/mleku/dendrite/pkg/ratio"
  24  )
  25  
  26  // Score holds the three-dimensional fitness evaluation.
  27  // All fields are exact rationals — no floating-point nondeterminism.
  28  type Score struct {
  29  	Source       ratio.Ratio `json:"source"`        // 0..1 structural AST similarity
  30  	Binary       ratio.Ratio `json:"binary"`        // 0..1 compiled binary similarity
  31  	Behav        ratio.Ratio `json:"behav"`         // 0..1 behavioral equivalence
  32  	Overall      ratio.Ratio `json:"overall"`       // weighted combination
  33  	CompileError string      `json:"compile_error,omitempty"`
  34  }
  35  
  36  // Compute sets the Overall score as a weighted combination.
  37  // Behavioral equivalence dominates — "does the code do the same thing?"
  38  // Weights: 3/20 source + 1/20 binary + 16/20 behavioral = 1.
  39  func (s *Score) Compute() {
  40  	s.Overall = ratio.New(3, 20).Mul(s.Source).
  41  		Add(ratio.New(1, 20).Mul(s.Binary)).
  42  		Add(ratio.New(16, 20).Mul(s.Behav))
  43  }
  44  
  45  // SourceSimilarity compares two Go source strings structurally.
  46  // It parses both into ASTs and measures the overlap of declarations:
  47  // package name, type names, function names, method names.
  48  func SourceSimilarity(original, emitted string) ratio.Ratio {
  49  	origDecls := extractDecls(original)
  50  	emitDecls := extractDecls(emitted)
  51  
  52  	if len(origDecls) == 0 {
  53  		return ratio.Zero
  54  	}
  55  
  56  	// Count how many original declarations appear in the emitted code.
  57  	matches := 0
  58  	for decl := range origDecls {
  59  		if emitDecls[decl] {
  60  			matches++
  61  		}
  62  	}
  63  
  64  	return ratio.New(int64(matches), int64(len(origDecls)))
  65  }
  66  
  67  // extractDecls parses Go source and returns a set of declaration signatures.
  68  func extractDecls(src string) map[string]bool {
  69  	decls := make(map[string]bool)
  70  
  71  	fset := token.NewFileSet()
  72  	f, err := parser.ParseFile(fset, "", src, parser.SkipObjectResolution)
  73  	if err != nil {
  74  		// If it doesn't parse, try to extract what we can from text.
  75  		return extractDeclsFromText(src)
  76  	}
  77  
  78  	if f.Name != nil {
  79  		decls["package:"+f.Name.Name] = true
  80  	}
  81  
  82  	for _, d := range f.Decls {
  83  		switch decl := d.(type) {
  84  		case *ast.FuncDecl:
  85  			if decl.Recv != nil && len(decl.Recv.List) > 0 {
  86  				// Method.
  87  				decls["method:"+decl.Name.Name] = true
  88  			} else {
  89  				decls["func:"+decl.Name.Name] = true
  90  			}
  91  		case *ast.GenDecl:
  92  			for _, spec := range decl.Specs {
  93  				switch s := spec.(type) {
  94  				case *ast.TypeSpec:
  95  					decls["type:"+s.Name.Name] = true
  96  				case *ast.ImportSpec:
  97  					if s.Path != nil {
  98  						decls["import:"+s.Path.Value] = true
  99  					}
 100  				}
 101  			}
 102  		}
 103  	}
 104  
 105  	return decls
 106  }
 107  
 108  // extractDeclsFromText does a rough text-based extraction for non-parseable source.
 109  func extractDeclsFromText(src string) map[string]bool {
 110  	decls := make(map[string]bool)
 111  	for _, line := range strings.Split(src, "\n") {
 112  		line = strings.TrimSpace(line)
 113  		if strings.HasPrefix(line, "package ") {
 114  			decls["package:"+strings.Fields(line)[1]] = true
 115  		}
 116  		if strings.HasPrefix(line, "func ") {
 117  			// Extract function name.
 118  			rest := strings.TrimPrefix(line, "func ")
 119  			if idx := strings.IndexByte(rest, '('); idx > 0 {
 120  				name := strings.TrimSpace(rest[:idx])
 121  				if strings.Contains(name, ")") {
 122  					// Method: "func (x T) Name("
 123  					parts := strings.SplitAfter(name, ")")
 124  					if len(parts) > 1 {
 125  						decls["method:"+strings.TrimSpace(parts[1])] = true
 126  					}
 127  				} else {
 128  					decls["func:"+name] = true
 129  				}
 130  			}
 131  		}
 132  		if strings.HasPrefix(line, "type ") {
 133  			fields := strings.Fields(line)
 134  			if len(fields) >= 2 {
 135  				decls["type:"+fields[1]] = true
 136  			}
 137  		}
 138  	}
 139  	return decls
 140  }
 141  
 142  // BinarySimilarity compares two compiled binaries.
 143  // Measures: size ratio and shared byte sequences.
 144  func BinarySimilarity(originalPath, emittedPath string) ratio.Ratio {
 145  	origData, err := os.ReadFile(originalPath)
 146  	if err != nil {
 147  		return ratio.Zero
 148  	}
 149  	emitData, err := os.ReadFile(emittedPath)
 150  	if err != nil {
 151  		return ratio.Zero
 152  	}
 153  
 154  	if len(origData) == 0 || len(emitData) == 0 {
 155  		return ratio.Zero
 156  	}
 157  
 158  	// Size ratio — min/max so it's always <= 1.
 159  	small, large := int64(len(emitData)), int64(len(origData))
 160  	if small > large {
 161  		small, large = large, small
 162  	}
 163  	sizeRatio := ratio.New(small, large)
 164  
 165  	// Shared 4-byte sequences (rough structural similarity).
 166  	// Sample to keep it fast.
 167  	const chunkSize = 4     // 2^2 byte chunks
 168  	const maxSamples = 25600 // 10^2 × 2^8 — epoch-aligned with binary chunk boundary
 169  
 170  	origChunks := make(map[string]bool)
 171  	step := intMax(1, (len(origData)-chunkSize)/maxSamples)
 172  	for i := 0; i+chunkSize <= len(origData); i += step {
 173  		origChunks[string(origData[i:i+chunkSize])] = true
 174  	}
 175  
 176  	shared := 0
 177  	total := 0
 178  	step = intMax(1, (len(emitData)-chunkSize)/maxSamples)
 179  	for i := 0; i+chunkSize <= len(emitData); i += step {
 180  		total++
 181  		if origChunks[string(emitData[i:i+chunkSize])] {
 182  			shared++
 183  		}
 184  	}
 185  
 186  	chunkSim := ratio.Zero
 187  	if total > 0 {
 188  		chunkSim = ratio.New(int64(shared), int64(total))
 189  	}
 190  
 191  	return ratio.Half.Mul(sizeRatio).Add(ratio.Half.Mul(chunkSim))
 192  }
 193  
 194  // BehavioralSimilarity runs both binaries with the same input and
 195  // compares their stdout output. This is the real fitness test.
 196  func BehavioralSimilarity(originalBin, emittedBin string, args []string, timeout string) ratio.Ratio {
 197  	origOut := runBinary(originalBin, args, timeout)
 198  	emitOut := runBinary(emittedBin, args, timeout)
 199  
 200  	if len(origOut) == 0 && len(emitOut) == 0 {
 201  		// Both produce no output — trivially equivalent.
 202  		return ratio.One
 203  	}
 204  	if len(origOut) == 0 || len(emitOut) == 0 {
 205  		return ratio.Zero
 206  	}
 207  
 208  	// Exact match.
 209  	if bytes.Equal(origOut, emitOut) {
 210  		return ratio.One
 211  	}
 212  
 213  	// Line-level similarity.
 214  	origLines := strings.Split(string(origOut), "\n")
 215  	emitLines := strings.Split(string(emitOut), "\n")
 216  
 217  	return lineSimilarity(origLines, emitLines)
 218  }
 219  
 220  // runBinary executes a binary in a temp directory and captures stdout.
 221  // Using a temp dir prevents the binary from clobbering files in the project.
 222  func runBinary(binPath string, args []string, timeout string) []byte {
 223  	if timeout == "" {
 224  		timeout = "10s"
 225  	}
 226  
 227  	tmpDir, err := os.MkdirTemp("", "fitness-run-*")
 228  	if err != nil {
 229  		return nil
 230  	}
 231  	defer os.RemoveAll(tmpDir)
 232  
 233  	cmd := exec.Command("timeout", append([]string{timeout, binPath}, args...)...)
 234  	cmd.Dir = tmpDir
 235  	cmd.Env = cleanGoEnv("/home/mleku/sdk/go1.24.6")
 236  	out, _ := cmd.CombinedOutput()
 237  	return out
 238  }
 239  
 240  // cleanGoEnv returns an environment with GOROOT, PATH, and GOTOOLCHAIN
 241  // set to use the specified Go root, preventing contamination from the
 242  // system Go installation.
 243  func cleanGoEnv(root string) []string {
 244  	env := os.Environ()
 245  	clean := make([]string, 0, len(env)+3)
 246  	for _, e := range env {
 247  		if strings.HasPrefix(e, "GOROOT=") ||
 248  			strings.HasPrefix(e, "GOTOOLCHAIN=") ||
 249  			strings.HasPrefix(e, "PATH=") {
 250  			continue
 251  		}
 252  		clean = append(clean, e)
 253  	}
 254  	clean = append(clean,
 255  		"GOROOT="+root,
 256  		"GOTOOLCHAIN=local",
 257  		"PATH="+filepath.Join(root, "bin")+":"+os.Getenv("PATH"),
 258  	)
 259  	return clean
 260  }
 261  
 262  // lineSimilarity computes how well the emitted output reproduces the original.
 263  // Uses a combination of exact line matches and token overlap to measure
 264  // partial matches (e.g., format strings with verbs stripped).
 265  func lineSimilarity(orig, emitted []string) ratio.Ratio {
 266  	if len(orig) == 0 {
 267  		return ratio.Zero
 268  	}
 269  
 270  	// First pass: exact line matches (highest confidence).
 271  	emitSet := make(map[string]bool, len(emitted))
 272  	for _, line := range emitted {
 273  		emitSet[strings.TrimSpace(line)] = true
 274  	}
 275  
 276  	exactMatches := 0
 277  	for _, line := range orig {
 278  		if emitSet[strings.TrimSpace(line)] {
 279  			exactMatches++
 280  		}
 281  	}
 282  
 283  	// Second pass: token overlap — do the emitted lines share words
 284  	// with the original? This catches format-string matches where
 285  	// "abiogenesis: nodes" partially matches "abiogenesis: 72 nodes".
 286  	emitTokens := make(map[string]bool)
 287  	for _, line := range emitted {
 288  		for _, tok := range strings.Fields(strings.TrimSpace(line)) {
 289  			if len(tok) >= 3 { // skip short tokens
 290  				emitTokens[strings.ToLower(tok)] = true
 291  			}
 292  		}
 293  	}
 294  
 295  	tokenScore := ratio.Zero
 296  	for _, line := range orig {
 297  		trimmed := strings.TrimSpace(line)
 298  		if trimmed == "" {
 299  			continue
 300  		}
 301  		toks := strings.Fields(trimmed)
 302  		if len(toks) == 0 {
 303  			continue
 304  		}
 305  		hits := 0
 306  		for _, tok := range toks {
 307  			if len(tok) >= 3 && emitTokens[strings.ToLower(tok)] {
 308  				hits++
 309  			}
 310  		}
 311  		total := 0
 312  		for _, tok := range toks {
 313  			if len(tok) >= 3 {
 314  				total++
 315  			}
 316  		}
 317  		if total > 0 {
 318  			tokenScore = tokenScore.Add(ratio.New(int64(hits), int64(total)))
 319  		}
 320  	}
 321  
 322  	nonEmpty := 0
 323  	for _, line := range orig {
 324  		if strings.TrimSpace(line) != "" {
 325  			nonEmpty++
 326  		}
 327  	}
 328  	if nonEmpty == 0 {
 329  		nonEmpty = 1
 330  	}
 331  	avgTokenSim := tokenScore.Div(ratio.FromInt(int64(nonEmpty)))
 332  
 333  	// Combine: exact matches are worth more, token overlap adds partial credit.
 334  	exactRat := ratio.New(int64(exactMatches), int64(len(orig)))
 335  	return ratio.Half.Mul(exactRat).Add(ratio.Half.Mul(avgTokenSim))
 336  }
 337  
 338  // CompileTo compiles a Go source file into a binary in the given directory.
 339  // Returns the path to the binary, or error.
 340  func CompileTo(sourceFile, outputBin, goRoot string) error {
 341  	dir := filepath.Dir(outputBin)
 342  	os.MkdirAll(dir, 0o755)
 343  
 344  	// Read the source to check for self-imports.
 345  	src, err := os.ReadFile(sourceFile)
 346  	if err != nil {
 347  		return err
 348  	}
 349  
 350  	// Write a go.mod. If the source imports dendrite packages, add
 351  	// require + replace directives so the build can resolve them.
 352  	modPath := filepath.Join(dir, "go.mod")
 353  	if _, err := os.Stat(modPath); os.IsNotExist(err) {
 354  		modContent := buildGoMod(string(src), dir)
 355  		os.WriteFile(modPath, []byte(modContent), 0o644)
 356  	}
 357  
 358  	// Copy source to dir/main.go.
 359  	mainPath := filepath.Join(dir, "main.go")
 360  	if err := os.WriteFile(mainPath, src, 0o644); err != nil {
 361  		return err
 362  	}
 363  
 364  	goBin := filepath.Join(goRoot, "bin", "go")
 365  	env := cleanGoEnv(goRoot)
 366  
 367  	// Run go mod tidy to reconcile dependencies before building.
 368  	tidy := exec.Command(goBin, "mod", "tidy")
 369  	tidy.Dir = dir
 370  	tidy.Env = env
 371  	tidy.CombinedOutput() // best-effort; build will report real errors
 372  
 373  	cmd := exec.Command(goBin, "build", "-o", outputBin, ".")
 374  	cmd.Dir = dir
 375  	cmd.Env = env
 376  	out, err := cmd.CombinedOutput()
 377  	if err != nil {
 378  		return &CompileError{Output: string(out), Err: err}
 379  	}
 380  
 381  	return nil
 382  }
 383  
 384  const selfModule = "git.mleku.dev/mleku/dendrite"
 385  
 386  // buildGoMod generates a go.mod for offspring compilation. If the source
 387  // imports dendrite sub-packages, the go.mod includes require/replace
 388  // directives pointing to the local source tree and copies go.sum for
 389  // transitive dependency resolution. buildDir is the directory where
 390  // go.mod and go.sum will be written.
 391  func buildGoMod(source string, buildDir string) string {
 392  	var b strings.Builder
 393  	b.WriteString("module offspring\n\ngo 1.24\n")
 394  
 395  	if !strings.Contains(source, selfModule) {
 396  		return b.String()
 397  	}
 398  
 399  	// Find the dendrite module root by walking up from cwd.
 400  	modRoot := findModuleRoot(selfModule)
 401  	if modRoot == "" {
 402  		return b.String()
 403  	}
 404  
 405  	// Read the parent go.mod to copy its require blocks.
 406  	parentMod, err := os.ReadFile(filepath.Join(modRoot, "go.mod"))
 407  	if err != nil {
 408  		return b.String()
 409  	}
 410  
 411  	// Extract require blocks from parent go.mod.
 412  	requires := extractRequireBlocks(string(parentMod))
 413  	if requires != "" {
 414  		b.WriteString("\n")
 415  		b.WriteString(requires)
 416  	}
 417  
 418  	// Add self-require + replace.
 419  	b.WriteString("\nrequire " + selfModule + " v0.0.0\n")
 420  	b.WriteString("\nreplace " + selfModule + " => " + modRoot + "\n")
 421  
 422  	// Copy go.sum for transitive dependency resolution.
 423  	sumSrc := filepath.Join(modRoot, "go.sum")
 424  	sumDst := filepath.Join(buildDir, "go.sum")
 425  	if sumData, err := os.ReadFile(sumSrc); err == nil {
 426  		os.WriteFile(sumDst, sumData, 0o644)
 427  	}
 428  
 429  	return b.String()
 430  }
 431  
 432  // findModuleRoot walks up from the current working directory looking for
 433  // a go.mod that declares the given module path.
 434  func findModuleRoot(modulePath string) string {
 435  	dir, err := os.Getwd()
 436  	if err != nil {
 437  		return ""
 438  	}
 439  	for {
 440  		modFile := filepath.Join(dir, "go.mod")
 441  		data, err := os.ReadFile(modFile)
 442  		if err == nil {
 443  			// Check if this go.mod declares our module.
 444  			for _, line := range strings.Split(string(data), "\n") {
 445  				line = strings.TrimSpace(line)
 446  				if strings.HasPrefix(line, "module ") {
 447  					mod := strings.TrimSpace(strings.TrimPrefix(line, "module"))
 448  					if mod == modulePath {
 449  						return dir
 450  					}
 451  				}
 452  			}
 453  		}
 454  		parent := filepath.Dir(dir)
 455  		if parent == dir {
 456  			break
 457  		}
 458  		dir = parent
 459  	}
 460  	return ""
 461  }
 462  
 463  // extractRequireBlocks extracts all require(...) blocks and single require
 464  // lines from a go.mod string, excluding the module declaration.
 465  func extractRequireBlocks(gomod string) string {
 466  	var b strings.Builder
 467  	lines := strings.Split(gomod, "\n")
 468  	inBlock := false
 469  	for _, line := range lines {
 470  		trimmed := strings.TrimSpace(line)
 471  		if trimmed == "require (" {
 472  			inBlock = true
 473  			b.WriteString(line + "\n")
 474  			continue
 475  		}
 476  		if inBlock {
 477  			b.WriteString(line + "\n")
 478  			if trimmed == ")" {
 479  				inBlock = false
 480  			}
 481  			continue
 482  		}
 483  		// Single-line require (but not "module" line).
 484  		if strings.HasPrefix(trimmed, "require ") && !strings.HasPrefix(trimmed, "require (") {
 485  			b.WriteString(line + "\n")
 486  		}
 487  	}
 488  	return b.String()
 489  }
 490  
 491  // CompileError wraps a compilation failure with the compiler output.
 492  type CompileError struct {
 493  	Output string
 494  	Err    error
 495  }
 496  
 497  func (e *CompileError) Error() string {
 498  	return e.Err.Error() + ": " + e.Output
 499  }
 500  
 501  // Evaluate runs the full fitness evaluation pipeline.
 502  func Evaluate(originalSource, emittedSource, originalBin, goRoot string, testArgs []string) Score {
 503  	var s Score
 504  
 505  	// 1. Source similarity.
 506  	origSrc, err := os.ReadFile(originalSource)
 507  	if err == nil {
 508  		emitSrc, err2 := os.ReadFile(emittedSource)
 509  		if err2 == nil {
 510  			s.Source = SourceSimilarity(string(origSrc), string(emitSrc))
 511  		}
 512  	}
 513  
 514  	// 2. Compile the emitted source and compare binaries.
 515  	absEmitted, _ := filepath.Abs(emittedSource)
 516  	emitBinDir := filepath.Join(filepath.Dir(absEmitted), "_fitness_build")
 517  	emitBin := filepath.Join(emitBinDir, "offspring")
 518  	absOrigBin, _ := filepath.Abs(originalBin)
 519  	defer os.RemoveAll(emitBinDir)
 520  
 521  	compileErr := CompileTo(absEmitted, emitBin, goRoot)
 522  	if compileErr != nil {
 523  		s.CompileError = compileErr.Error()
 524  	} else {
 525  		s.Binary = BinarySimilarity(absOrigBin, emitBin)
 526  		s.Behav = BehavioralSimilarity(absOrigBin, emitBin, testArgs, "15s")
 527  	}
 528  
 529  	s.Compute()
 530  	return s
 531  }
 532  
 533  // intMax returns the larger of two ints.
 534  func intMax(a, b int) int {
 535  	if a > b {
 536  		return a
 537  	}
 538  	return b
 539  }
 540