Prune unused imports before rendering files

sszgen/core/file.go

@@ -8,9 +8,9 @@ import (
 	"text/template"
 )
 
-// DefaultSSZImports is seeded into the ImportNamer for a generated SSZ file. The
-// entries are always emitted (ImportPairs serializes the whole set), which is
-// valid only when a complete method set is generated into one file.
+// DefaultSSZImports is seeded into the ImportNamer for a generated SSZ file.
+// Seeding claims the identifiers up front so alias assignment is stable, but an
+// entry is only emitted if the rendered body actually references it.
 var DefaultSSZImports = map[string]string{
 	"github.com/OffchainLabs/methodical-ssz/ssz": "ssz",
 	"fmt":             "",
@@ -41,6 +41,10 @@ func RenderFile(pkgName string, namer *ImportNamer, blockGroups ...[]string) ([]
 	for _, g := range blockGroups {
 		blocks = append(blocks, g...)
 	}
+
+	body := strings.Join(blocks, "\n\n")
+	namer.PruneUnreferenced(body)
+
 	buf := bytes.NewBuffer(nil)
 	err := fileTemplate.Execute(buf, struct {
 		Package string
@@ -49,7 +53,7 @@ func RenderFile(pkgName string, namer *ImportNamer, blockGroups ...[]string) ([]
 	}{
 		Package: pkgName,
 		Imports: namer.ImportPairs(),
-		Blocks:  strings.Join(blocks, "\n\n"),
+		Blocks:  body,
 	})
 	if err != nil {
 		return nil, err

sszgen/core/imports.go

@@ -2,6 +2,8 @@ package core
 
 import (
 	"fmt"
+	"go/scanner"
+	"go/token"
 	"go/types"
 	"sort"
 	"strings"
@@ -35,10 +37,11 @@ type importNode struct {
 	depth    int
 	imported bool
 	// used marks packages whose name has been handed to code generation
-	// (NameString) — only those are emitted by ImportPairs, so reserving more
-	// packages than the generated code references is harmless. A used node's
-	// name is frozen: it may already appear in generated text, so the
-	// shortest-path-wins reorganization never steals from it.
+	// (NameString) or seeded as a default — only those are emitted by
+	// ImportPairs, so reserving more packages than the generated code
+	// references is harmless. A used node's name is frozen: it may already
+	// appear in generated text, so the shortest-path-wins reorganization never
+	// steals from it.
 	used bool
 	// name is the identifier generated code uses for this package. bare means
 	// the import is emitted without an explicit alias (the identifier is the
@@ -214,6 +217,36 @@ func (n *ImportNamer) ImportPairs() string {
 	return strings.Join(lines, "\n")
 }
 
+// PruneUnreferenced drops every registered import whose identifier does not
+// appear in src, the rendered body of the file being generated.
+func (n *ImportNamer) PruneUnreferenced(src string) {
+	idents := referencedIdents(src)
+	n.walk(func(nd *importNode) {
+		if !idents[nd.name] {
+			nd.used = false
+		}
+	})
+}
+
+// referencedIdents returns the identifiers appearing as Go tokens in src.
+func referencedIdents(src string) map[string]bool {
+	fset := token.NewFileSet()
+
+	var s scanner.Scanner
+	s.Init(fset.AddFile("", fset.Base(), len(src)), []byte(src), nil, 0)
+
+	idents := make(map[string]bool)
+	for {
+		_, tok, lit := s.Scan()
+		switch tok {
+		case token.EOF:
+			return idents
+		case token.IDENT:
+			idents[lit] = true
+		}
+	}
+}
+
 // ImportSource renders a complete import declaration block (see ImportPairs).
 func (n *ImportNamer) ImportSource() string {
 	return fmt.Sprintf("import (\n%s\n)\n", n.ImportPairs())

sszgen/core/imports_test.go

@@ -107,12 +107,47 @@ func TestImportNamerReserveUnusedNotEmitted(t *testing.T) {
 	if strings.Contains(pairs, "widget") {
 		t.Fatalf("reserved-but-unused package emitted:\n%s", pairs)
 	}
-	// defaults are always emitted (body templates reference them unconditionally)
+	// seeded defaults are emitted until PruneUnreferenced says otherwise
 	if !strings.Contains(pairs, `"fmt"`) {
 		t.Fatalf("default missing:\n%s", pairs)
 	}
 }
 
+func TestImportNamerPruneUnreferenced(t *testing.T) {
+	n := NewImportNamer("github.com/example/target", DefaultSSZImports)
+	n.NameString("github.com/example/engine/v1")
+	n.PruneUnreferenced(`func (c *X) MarshalSSZ() ([]byte, error) {
+	if c.Inner == nil {
+		c.Inner = new(v1.Inner)
+	}
+	return nil, fmt.Errorf("nope")
+}`)
+	pairs := n.ImportPairs()
+	for _, want := range []string{`"fmt"`, `v1 "github.com/example/engine/v1"`} {
+		if !strings.Contains(pairs, want) {
+			t.Fatalf("referenced import %s dropped:\n%s", want, pairs)
+		}
+	}
+	for _, unwanted := range []string{"encoding/binary", "methodical-ssz/ssz"} {
+		if strings.Contains(pairs, unwanted) {
+			t.Fatalf("unreferenced import %s emitted:\n%s", unwanted, pairs)
+		}
+	}
+}
+
+func TestImportNamerPruneIgnoresCommentsAndStrings(t *testing.T) {
+	n := NewImportNamer("github.com/example/target", DefaultSSZImports)
+	n.PruneUnreferenced(`// binary.LittleEndian would go here
+func (c *X) Err() error { return fmt.Errorf("ssz.ErrBytesLength: %w", nil) }`)
+	pairs := n.ImportPairs()
+	if !strings.Contains(pairs, `"fmt"`) {
+		t.Fatalf("referenced import dropped:\n%s", pairs)
+	}
+	if strings.Contains(pairs, "encoding/binary") || strings.Contains(pairs, "methodical-ssz/ssz") {
+		t.Fatalf("comment/string mention counted as a reference:\n%s", pairs)
+	}
+}
+
 // Shortest path wins regardless of arrival order: a shorter path arriving
 // later steals the bare name from a reserved (not yet used) longer path, which
 // is renamed via its ancestor segments.

sszgen/render/renderer_test.go

@@ -2,6 +2,9 @@ package render
 
 import (
 	"flag"
+	"go/ast"
+	"go/parser"
+	"go/token"
 	"os"
 	"path/filepath"
 	"strings"
@@ -135,6 +138,77 @@ func TestRenderTopLevelOverlaySkipped(t *testing.T) {
 	}
 }
 
+func TestRenderOmitsUnusedImports(t *testing.T) {
+	inner := &gentypes.ValueContainer{
+		Name:    "Inner",
+		Package: "github.com/example/foo",
+		Contents: []gentypes.ContainerField{
+			{Key: "X", Value: &gentypes.ValueList{MaxSize: 64, ElementValue: &gentypes.ValueByte{Name: "byte"}}},
+		},
+	}
+	vc := &gentypes.ValueContainer{
+		Name:    "NoIntegers",
+		Package: "github.com/example/foo",
+		Contents: []gentypes.ContainerField{
+			{Key: "A", Value: &gentypes.ValueVector{Size: 32, ElementValue: &gentypes.ValueByte{Name: "byte"}}},
+			{Key: "B", Value: &gentypes.ValueBool{Name: "bool"}},
+			{Key: "C", Value: &gentypes.ValueList{MaxSize: 256, ElementValue: &gentypes.ValueByte{Name: "byte"}}},
+			{Key: "D", Value: &gentypes.ValuePointer{Referent: inner}},
+		},
+	}
+	got, err := Render("github.com/example/foo", "", []gentypes.ValRep{vc})
+	if err != nil {
+		t.Fatalf("render.Render: %v", err)
+	}
+	if strings.Contains(string(got), "encoding/binary") {
+		t.Fatalf("no field uses binary.LittleEndian, but encoding/binary was imported:\n%s", got)
+	}
+	assertNoUnusedImports(t, got)
+	// the imports it does need are still there
+	assertNoUnusedImports(t, mustRender(t, exampleContainer()))
+}
+
+func mustRender(t *testing.T, vr gentypes.ValRep) []byte {
+	t.Helper()
+	got, err := Render("github.com/example/foo", "", []gentypes.ValRep{vr})
+	if err != nil {
+		t.Fatalf("render.Render: %v", err)
+	}
+	return got
+}
+
+// assertNoUnusedImports fails if src imports a package whose identifier appears
+// nowhere outside the import block — the compile error this all guards against.
+func assertNoUnusedImports(t *testing.T, src []byte) {
+	t.Helper()
+	f, err := parser.ParseFile(token.NewFileSet(), "generated.ssz.go", src, 0)
+	if err != nil {
+		t.Fatalf("generated file does not parse: %v", err)
+	}
+	referenced := make(map[string]bool)
+	for _, d := range f.Decls {
+		if g, ok := d.(*ast.GenDecl); ok && g.Tok == token.IMPORT {
+			continue
+		}
+		ast.Inspect(d, func(n ast.Node) bool {
+			if id, ok := n.(*ast.Ident); ok {
+				referenced[id.Name] = true
+			}
+			return true
+		})
+	}
+	for _, spec := range f.Imports {
+		path := strings.Trim(spec.Path.Value, `"`)
+		name := path[strings.LastIndex(path, "/")+1:]
+		if spec.Name != nil {
+			name = spec.Name.Name
+		}
+		if !referenced[name] {
+			t.Fatalf("import %q is not referenced by the generated code", path)
+		}
+	}
+}
+
 // Render requires a package path (or an override) before doing any work.
 func TestRenderEmptyPackagePath(t *testing.T) {
 	if _, err := Render("", "", nil); err == nil {

sszgen/testdata/delegate.ssz.go.expected

@@ -1,7 +1,6 @@
 package delegatefixture
 
 import (
-	binary "encoding/binary"
 	"fmt"
 	ssz "github.com/OffchainLabs/methodical-ssz/ssz"
 )

sszgen/testdata/uint256.ssz.go.fixture

@@ -1,7 +1,6 @@
 package faketypes
 
 import (
-	binary "encoding/binary"
 	"fmt"
 	ssz "github.com/OffchainLabs/methodical-ssz/ssz"
 )