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