loadspec.go

     1  // Copyright 2026 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package specgen
     6  
     7  import (
     8  	"fmt"
     9  	"go/ast"
    10  	"go/build"
    11  	"go/importer"
    12  	"go/parser"
    13  	"go/token"
    14  	"go/types"
    15  	"path/filepath"
    16  	"simd/archsimd/_gen/specgen/specexpr"
    17  	"strings"
    18  )
    19  
    20  // specPackage represents the parsed _gen/spec package.
    21  type specPackage struct {
    22  	Fset      *token.FileSet
    23  	Pkg       *types.Package
    24  	TypesInfo *types.Info
    25  	Funcs     []*specFunc
    26  
    27  	TypeElems  map[types.Type]specexpr.Basic
    28  	TypeWidths map[types.Type]specexpr.Num
    29  
    30  	ElemTypes  map[specexpr.Basic]types.Type
    31  	WidthTypes map[specexpr.Num]types.Type
    32  
    33  	VecType   types.Type // Uninstantiated Vec type
    34  	ArrayType types.Type // Uninstantiated Array type
    35  	UintNType types.Type // Uninstantiated UintN type
    36  }
    37  
    38  // specFunc represents an exported function in the spec source.
    39  type specFunc struct {
    40  	Pkg          *specPackage
    41  	Name         string       // Source name
    42  	NameTmpl     specTemplate // API name template from `//specgen:name` directive, or same as Name.
    43  	Pos          token.Pos
    44  	Doc          specTemplate
    45  	Category     string
    46  	Commutative  bool
    47  	Sig          *types.Signature
    48  	TypeParams   []*types.TypeParam
    49  	Params       []*types.Var
    50  	Results      []*types.Var
    51  	Requirements []specexpr.Expr
    52  }
    53  
    54  // specTemplate is a template string, with placeholders of the form `{var}`,
    55  // which will be replaced with variable values from the solver.
    56  type specTemplate struct {
    57  	tmpl   string   // raw template string including patterns
    58  	fields [][2]int // start:end ranges of fields, including '{}'s, in ascending order
    59  }
    60  
    61  // specGoVersion is the oldest Go toolchain version that must be able to parse
    62  // and type-check the spec package.
    63  const specGoVersion = "go1.26"
    64  
    65  // loadAndTypeCheck imports and parses the spec package in dir, ensuring it is
    66  // target-independent, and type-checks it.
    67  func loadAndTypeCheck(ctx context, dir string) (*types.Package, *types.Info, []*ast.File) {
    68  	bp, err := build.ImportDir(dir, 0)
    69  	if err != nil {
    70  		ctx.errorf("failed to import spec directory %s: %s", dir, err)
    71  		return nil, nil, nil
    72  	}
    73  	if len(bp.AllTags) > 0 {
    74  		ctx.errorf("internal/spec must be target-independent, but found build tags: %v", bp.AllTags)
    75  		return nil, nil, nil
    76  	}
    77  
    78  	var astFiles []*ast.File
    79  	for _, name := range bp.GoFiles {
    80  		filePath := filepath.Join(bp.Dir, name)
    81  		file, err := parser.ParseFile(&ctx.root.fset, filePath, nil, parser.ParseComments)
    82  		if err != nil {
    83  			ctx.errorf("failed to parse %s: %s", filePath, err)
    84  			continue
    85  		}
    86  		astFiles = append(astFiles, file)
    87  	}
    88  
    89  	if len(astFiles) == 0 {
    90  		ctx.errorf("no Go source files found in directory %s", dir)
    91  		return nil, nil, nil
    92  	}
    93  
    94  	info := &types.Info{
    95  		Types:      make(map[ast.Expr]types.TypeAndValue),
    96  		Defs:       make(map[*ast.Ident]types.Object),
    97  		Uses:       make(map[*ast.Ident]types.Object),
    98  		Implicits:  make(map[ast.Node]types.Object),
    99  		Selections: make(map[*ast.SelectorExpr]*types.Selection),
   100  		Scopes:     make(map[ast.Node]*types.Scope),
   101  		Instances:  make(map[*ast.Ident]types.Instance),
   102  	}
   103  
   104  	// Note: Imported packages resolve against the running toolchain's standard
   105  	// library, not the dev tree's. For math and math/bits that is immaterial.
   106  	// Constraint: internal/spec may import only long-stable standard library packages.
   107  	conf := types.Config{
   108  		GoVersion: specGoVersion,
   109  		Importer:  importer.ForCompiler(&ctx.root.fset, "source", nil),
   110  		Error: func(err error) {
   111  			ctx.errorf("%s", err)
   112  		},
   113  	}
   114  
   115  	typesPkg, err := conf.Check("simd/internal/spec", &ctx.root.fset, astFiles, info)
   116  	if err != nil && typesPkg == nil {
   117  		return nil, nil, nil
   118  	}
   119  	if len(ctx.root.errors) > 0 {
   120  		return nil, nil, nil
   121  	}
   122  
   123  	return typesPkg, info, astFiles
   124  }
   125  
   126  // loadSpecPackage parses the spec package in the given directory path.
   127  func loadSpecPackage(ctx context, dir string, opts *LoadOptions) *specPackage {
   128  	typesPkg, info, astFiles := loadAndTypeCheck(ctx, dir)
   129  	if typesPkg == nil {
   130  		return nil
   131  	}
   132  
   133  	var pkg specPackage
   134  
   135  	// Gather exported functions
   136  	var funcs []*specFunc
   137  	for _, file := range astFiles {
   138  		// Gather directives
   139  		var category string
   140  		directives := make(map[*ast.Comment]ast.Directive)
   141  		for _, cg := range file.Comments {
   142  			for _, comment := range cg.List {
   143  				if dir, ok := ast.ParseDirective(comment.Slash, comment.Text); ok && dir.Tool == "specgen" {
   144  					switch dir.Name {
   145  					case "category":
   146  						// File-level directive
   147  						if category != "" {
   148  							ctx.at(dir.Pos()).errorf("multiple category directives in file")
   149  						}
   150  						category = dir.Args
   151  					case "name", "commutative", "require":
   152  						// Gather other directives to process with decls
   153  						directives[comment] = dir
   154  					default:
   155  						ctx.at(dir.Pos()).errorf("unknown //specgen directive")
   156  					}
   157  				}
   158  			}
   159  		}
   160  
   161  		for _, decl := range file.Decls {
   162  			d, ok := decl.(*ast.FuncDecl)
   163  			if !ok || !d.Name.IsExported() {
   164  				continue
   165  			}
   166  
   167  			obj := typesPkg.Scope().Lookup(d.Name.Name)
   168  			if obj == nil {
   169  				continue
   170  			}
   171  			fn, ok := obj.(*types.Func)
   172  			if !ok {
   173  				continue
   174  			}
   175  
   176  			sig := fn.Type().(*types.Signature)
   177  
   178  			var typeParams []*types.TypeParam
   179  			tparams := sig.TypeParams()
   180  			for tparam := range tparams.TypeParams() {
   181  				typeParams = append(typeParams, tparam)
   182  			}
   183  
   184  			var params []*types.Var
   185  			p := sig.Params()
   186  			for v := range p.Variables() {
   187  				params = append(params, v)
   188  			}
   189  
   190  			var results []*types.Var
   191  			r := sig.Results()
   192  			for v := range r.Variables() {
   193  				results = append(results, v)
   194  			}
   195  
   196  			f := &specFunc{
   197  				Pkg:        &pkg,
   198  				Name:       d.Name.Name,
   199  				Pos:        decl.Pos(),
   200  				Sig:        sig,
   201  				TypeParams: typeParams,
   202  				Params:     params,
   203  				Results:    results,
   204  				Category:   category,
   205  			}
   206  			f.NameTmpl = specTemplate{tmpl: f.Name}
   207  			if d.Doc != nil {
   208  				var err error
   209  				f.Doc, err = newSpecTemplate(d.Doc.Text())
   210  				if err != nil {
   211  					ctx.at(d.Doc.Pos()).errorf("malformed doc comment: %s", err)
   212  				}
   213  				for _, comment := range d.Doc.List {
   214  					if dir, ok := directives[comment]; ok {
   215  						delete(directives, comment)
   216  						switch dir.Name {
   217  						default:
   218  							panic("directive lists out of sync")
   219  						case "name":
   220  							f.NameTmpl, err = newSpecTemplate(dir.Args)
   221  							if err != nil {
   222  								ctx.at(dir.Pos()).errorf("malformed //specgen:name directive: %s", err)
   223  							}
   224  						case "commutative":
   225  							if dir.Args != "" {
   226  								ctx.at(dir.Pos()).errorf("malformed //specgen:commutative directive: expected no argument")
   227  							}
   228  							f.Commutative = true
   229  						case "require":
   230  							args, err := dir.ParseArgs()
   231  							if err != nil {
   232  								ctx.at(dir.Pos()).errorf("malformed //specgen:require directive: %s", err)
   233  								break
   234  							}
   235  							for _, arg := range args {
   236  								expr, err := specexpr.ParseExpr(arg.Arg)
   237  								if err != nil {
   238  									ctx.at(arg.Pos).errorf("failed to parse require argument %q: %s", arg.Arg, err)
   239  									continue
   240  								}
   241  								f.Requirements = append(f.Requirements, expr)
   242  							}
   243  						}
   244  					}
   245  				}
   246  			}
   247  
   248  			if opts.Filter != nil && !opts.Filter(d) {
   249  				continue
   250  			}
   251  
   252  			funcs = append(funcs, f)
   253  		}
   254  
   255  		for _, dir := range directives {
   256  			ctx.at(dir.Pos()).errorf("//%s:%s directive must be attached to a function", dir.Tool, dir.Name)
   257  		}
   258  	}
   259  
   260  	lookupType := func(name string) types.Type {
   261  		obj := typesPkg.Scope().Lookup(name)
   262  		if obj == nil {
   263  			ctx.errorf("type %q missing from package %s", name, typesPkg.Path())
   264  			return nil
   265  		}
   266  		tn, ok := obj.(*types.TypeName)
   267  		if !ok {
   268  			ctx.at(obj.Pos()).errorf("%s expected to be a type", obj.String())
   269  			return nil
   270  		}
   271  		return tn.Type()
   272  	}
   273  
   274  	// Gather types corresponding to shape constraints
   275  	typeElems := make(map[types.Type]specexpr.Basic)
   276  	elemTypes := make(map[specexpr.Basic]types.Type)
   277  	if eltOrMask := lookupType("EltOrMask"); eltOrMask != nil {
   278  		for _, elt := range typeSet(eltOrMask) {
   279  			basic := shapeElemType(elt)
   280  			typeElems[elt] = basic
   281  			elemTypes[basic] = elt
   282  		}
   283  	}
   284  	typeWidths := make(map[types.Type]specexpr.Num)
   285  	widthTypes := make(map[specexpr.Num]types.Type)
   286  	if width := lookupType("Width"); width != nil {
   287  		for _, width := range typeSet(width) {
   288  			val := shapeWidthVal(width)
   289  			typeWidths[width] = val
   290  			widthTypes[val] = width
   291  		}
   292  	}
   293  
   294  	// Gather other known types
   295  	vecType := lookupType("Vec")
   296  	arrayType := lookupType("Array")
   297  	uintNType := lookupType("UintN")
   298  
   299  	pkg = specPackage{
   300  		Fset:       &ctx.root.fset,
   301  		Pkg:        typesPkg,
   302  		TypesInfo:  info,
   303  		Funcs:      funcs,
   304  		TypeElems:  typeElems,
   305  		TypeWidths: typeWidths,
   306  		ElemTypes:  elemTypes,
   307  		WidthTypes: widthTypes,
   308  		VecType:    vecType,
   309  		ArrayType:  arrayType,
   310  		UintNType:  uintNType,
   311  	}
   312  	return &pkg
   313  }
   314  
   315  // newSpecTemplate parses spec template.
   316  func newSpecTemplate(tmpl string) (specTemplate, error) {
   317  	if !strings.ContainsAny(tmpl, "{}") {
   318  		return specTemplate{tmpl, nil}, nil
   319  	}
   320  
   321  	var fields [][2]int
   322  	for i := 0; i < len(tmpl); i++ {
   323  		switch tmpl[i] {
   324  		case '{':
   325  			j := i + strings.IndexByte(tmpl[i:], '}') + 1
   326  			if j <= i {
   327  				return specTemplate{}, fmt.Errorf("unclosed '{' in template %q", tmpl)
   328  			}
   329  			fields = append(fields, [2]int{i, j})
   330  			i = j - 1
   331  		case '}':
   332  			return specTemplate{}, fmt.Errorf("unmatched '}' in template %q", tmpl)
   333  		}
   334  	}
   335  	return specTemplate{
   336  		tmpl:   tmpl,
   337  		fields: fields,
   338  	}, nil
   339  }
   340  
   341  // expand replaces placeholders in template s by calling the lookup function to
   342  // resolve their values.
   343  func (s *specTemplate) expand(lookup func(string) string) string {
   344  	if len(s.fields) == 0 {
   345  		return s.tmpl
   346  	}
   347  	var buf strings.Builder
   348  	pos := 0
   349  	for _, field := range s.fields {
   350  		buf.WriteString(s.tmpl[pos:field[0]])
   351  		val := lookup(s.tmpl[field[0]+1 : field[1]-1])
   352  		buf.WriteString(val)
   353  		pos = field[1]
   354  	}
   355  	buf.WriteString(s.tmpl[pos:])
   356  	return buf.String()
   357  }
   358
View as plain text