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/token"
    11  	"go/types"
    12  	"simd/archsimd/_gen/specgen/specexpr"
    13  	"strings"
    14  
    15  	"golang.org/x/tools/go/packages"
    16  )
    17  
    18  // specPackage represents the parsed _gen/spec package.
    19  type specPackage struct {
    20  	Fset      *token.FileSet
    21  	Pkg       *types.Package
    22  	TypesInfo *types.Info
    23  	Funcs     []*specFunc
    24  
    25  	TypeElems  map[types.Type]specexpr.Basic
    26  	TypeWidths map[types.Type]specexpr.Num
    27  
    28  	ElemTypes  map[specexpr.Basic]types.Type
    29  	WidthTypes map[specexpr.Num]types.Type
    30  
    31  	VecType   types.Type // Uninstantiated Vec type
    32  	ArrayType types.Type // Uninstantiated Array type
    33  	UintNType types.Type // Uninstantiated UintN type
    34  }
    35  
    36  // specFunc represents an exported function in the spec source.
    37  type specFunc struct {
    38  	Pkg          *specPackage
    39  	Name         string       // Source name
    40  	NameTmpl     specTemplate // API name template from `//specgen:name` directive, or same as Name.
    41  	Pos          token.Pos
    42  	Doc          specTemplate
    43  	Sig          *types.Signature
    44  	TypeParams   []*types.TypeParam
    45  	Params       []*types.Var
    46  	Results      []*types.Var
    47  	Requirements []specexpr.Expr
    48  }
    49  
    50  // specTemplate is a template string, with placeholders of the form `{var}`,
    51  // which will be replaced with variable values from the solver.
    52  type specTemplate struct {
    53  	tmpl   string   // raw template string including patterns
    54  	fields [][2]int // start:end ranges of fields, including '{}'s, in ascending order
    55  }
    56  
    57  // loadSpecPackage parses the spec package in the given directory path.
    58  func loadSpecPackage(ctx context, dir string, opts *LoadOptions) *specPackage {
    59  	cfg := &packages.Config{
    60  		Mode: packages.LoadSyntax,
    61  		Dir:  dir,
    62  		Fset: &ctx.root.fset,
    63  	}
    64  
    65  	pkgs, err := packages.Load(cfg, ".")
    66  	if err != nil {
    67  		ctx.errorf("failed to load package: %s", err)
    68  		return nil
    69  	}
    70  	if len(pkgs) == 0 {
    71  		ctx.errorf("no package found in directory %s", dir)
    72  		return nil
    73  	}
    74  	if len(pkgs[0].Errors) > 0 {
    75  		for _, err := range pkgs[0].Errors {
    76  			ctx.errorf("%s", err)
    77  		}
    78  		return nil
    79  	}
    80  
    81  	srcPkg := pkgs[0]
    82  	fset := srcPkg.Fset
    83  	info := srcPkg.TypesInfo
    84  
    85  	var pkg specPackage
    86  
    87  	// Gather exported functions
    88  	var funcs []*specFunc
    89  	for _, file := range srcPkg.Syntax {
    90  		for _, decl := range file.Decls {
    91  			d, ok := decl.(*ast.FuncDecl)
    92  			if !ok || !d.Name.IsExported() {
    93  				continue
    94  			}
    95  			if opts.Filter != nil && !opts.Filter(d) {
    96  				continue
    97  			}
    98  
    99  			obj := srcPkg.Types.Scope().Lookup(d.Name.Name)
   100  			if obj == nil {
   101  				continue
   102  			}
   103  			fn, ok := obj.(*types.Func)
   104  			if !ok {
   105  				continue
   106  			}
   107  
   108  			sig := fn.Type().(*types.Signature)
   109  
   110  			var typeParams []*types.TypeParam
   111  			tparams := sig.TypeParams()
   112  			for tparam := range tparams.TypeParams() {
   113  				typeParams = append(typeParams, tparam)
   114  			}
   115  
   116  			var params []*types.Var
   117  			p := sig.Params()
   118  			for v := range p.Variables() {
   119  				params = append(params, v)
   120  			}
   121  
   122  			var results []*types.Var
   123  			r := sig.Results()
   124  			for v := range r.Variables() {
   125  				results = append(results, v)
   126  			}
   127  
   128  			f := &specFunc{
   129  				Pkg:        &pkg,
   130  				Name:       d.Name.Name,
   131  				Pos:        decl.Pos(),
   132  				Sig:        sig,
   133  				TypeParams: typeParams,
   134  				Params:     params,
   135  				Results:    results,
   136  			}
   137  			f.NameTmpl = specTemplate{tmpl: f.Name}
   138  			if d.Doc != nil {
   139  				f.Doc, err = newSpecTemplate(d.Doc.Text())
   140  				if err != nil {
   141  					ctx.at(d.Doc.Pos()).errorf("malformed doc comment: %s", err)
   142  				}
   143  				for _, comment := range d.Doc.List {
   144  					if dir, ok := ast.ParseDirective(comment.Slash, comment.Text); ok && dir.Tool == "specgen" {
   145  						switch dir.Name {
   146  						default:
   147  							ctx.at(dir.Pos()).errorf("unknown //specgen directive")
   148  						case "name":
   149  							f.NameTmpl, err = newSpecTemplate(dir.Args)
   150  							if err != nil {
   151  								ctx.at(dir.Pos()).errorf("malformed //specgen:name directive: %s", err)
   152  							}
   153  						case "require":
   154  							args, err := dir.ParseArgs()
   155  							if err != nil {
   156  								ctx.at(dir.Pos()).errorf("malformed //specgen:require directive: %s", err)
   157  								break
   158  							}
   159  							for _, arg := range args {
   160  								expr, err := specexpr.ParseExpr(arg.Arg)
   161  								if err != nil {
   162  									ctx.at(arg.Pos).errorf("failed to parse require argument %q: %s", arg.Arg, err)
   163  									continue
   164  								}
   165  								f.Requirements = append(f.Requirements, expr)
   166  							}
   167  						}
   168  					}
   169  				}
   170  			}
   171  
   172  			funcs = append(funcs, f)
   173  		}
   174  	}
   175  
   176  	lookupType := func(name string) types.Type {
   177  		obj := srcPkg.Types.Scope().Lookup(name)
   178  		if obj == nil {
   179  			ctx.errorf("type %q missing from package %s", name, srcPkg.PkgPath)
   180  			return nil
   181  		}
   182  		tn, ok := obj.(*types.TypeName)
   183  		if !ok {
   184  			ctx.at(obj.Pos()).errorf("%s expected to be a type", obj.String())
   185  			return nil
   186  		}
   187  		return tn.Type()
   188  	}
   189  
   190  	// Gather types corresponding to shape constraints
   191  	typeElems := make(map[types.Type]specexpr.Basic)
   192  	elemTypes := make(map[specexpr.Basic]types.Type)
   193  	if eltOrMask := lookupType("EltOrMask"); eltOrMask != nil {
   194  		for _, elt := range typeSet(eltOrMask) {
   195  			basic := shapeElemType(elt)
   196  			typeElems[elt] = basic
   197  			elemTypes[basic] = elt
   198  		}
   199  	}
   200  	typeWidths := make(map[types.Type]specexpr.Num)
   201  	widthTypes := make(map[specexpr.Num]types.Type)
   202  	if width := lookupType("Width"); width != nil {
   203  		for _, width := range typeSet(width) {
   204  			val := shapeWidthVal(width)
   205  			typeWidths[width] = val
   206  			widthTypes[val] = width
   207  		}
   208  	}
   209  
   210  	// Gather other known types
   211  	vecType := lookupType("Vec")
   212  	arrayType := lookupType("Array")
   213  	uintNType := lookupType("UintN")
   214  
   215  	pkg = specPackage{
   216  		Fset:       fset,
   217  		Pkg:        srcPkg.Types,
   218  		TypesInfo:  info,
   219  		Funcs:      funcs,
   220  		TypeElems:  typeElems,
   221  		TypeWidths: typeWidths,
   222  		ElemTypes:  elemTypes,
   223  		WidthTypes: widthTypes,
   224  		VecType:    vecType,
   225  		ArrayType:  arrayType,
   226  		UintNType:  uintNType,
   227  	}
   228  	return &pkg
   229  }
   230  
   231  // newSpecTemplate parses spec template.
   232  func newSpecTemplate(tmpl string) (specTemplate, error) {
   233  	if !strings.ContainsAny(tmpl, "{}") {
   234  		return specTemplate{tmpl, nil}, nil
   235  	}
   236  
   237  	var fields [][2]int
   238  	for i := 0; i < len(tmpl); i++ {
   239  		switch tmpl[i] {
   240  		case '{':
   241  			j := i + strings.IndexByte(tmpl[i:], '}') + 1
   242  			if j <= i {
   243  				return specTemplate{}, fmt.Errorf("unclosed '{' in template %q", tmpl)
   244  			}
   245  			fields = append(fields, [2]int{i, j})
   246  			i = j - 1
   247  		case '}':
   248  			return specTemplate{}, fmt.Errorf("unmatched '}' in template %q", tmpl)
   249  		}
   250  	}
   251  	return specTemplate{
   252  		tmpl:   tmpl,
   253  		fields: fields,
   254  	}, nil
   255  }
   256  
   257  // expand replaces placeholders in template s by calling the lookup function to
   258  // resolve their values.
   259  func (s *specTemplate) expand(lookup func(string) string) string {
   260  	if len(s.fields) == 0 {
   261  		return s.tmpl
   262  	}
   263  	var buf strings.Builder
   264  	pos := 0
   265  	for _, field := range s.fields {
   266  		buf.WriteString(s.tmpl[pos:field[0]])
   267  		val := lookup(s.tmpl[field[0]+1 : field[1]-1])
   268  		buf.WriteString(val)
   269  		pos = field[1]
   270  	}
   271  	buf.WriteString(s.tmpl[pos:])
   272  	return buf.String()
   273  }
   274
View as plain text