expand.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/types"
    10  	"regexp"
    11  	"simd/archsimd/_gen/specgen/specexpr"
    12  	"strings"
    13  )
    14  
    15  func (sFn *specFunc) expand(ctx context, opts *LoadOptions) []*Func {
    16  	ctx = ctx.at(sFn.Pos)
    17  
    18  	var solver specexpr.Solver
    19  
    20  	if opts.Trace != nil {
    21  		fmt.Fprintf(opts.Trace, "## %s%s\n", sFn.Name, sFn.Sig)
    22  		solver.SetTrace(opts.Trace)
    23  	}
    24  
    25  	// Declare domains of type parameters
    26  	typeParamVars := make(map[*types.TypeParam]specexpr.Variable)
    27  	for _, param := range sFn.TypeParams {
    28  		types, err := constraintToDomain(sFn.Pkg, param.Constraint())
    29  		if err != nil {
    30  			panic(err)
    31  		}
    32  		varDef := specexpr.Variable("$" + param.String())
    33  		solver.Declare(varDef, types)
    34  		typeParamVars[param] = varDef
    35  	}
    36  	// Bind shapes of all function parameters and results of Vec type
    37  	b := &argBinder{ctx, sFn.Pkg, &solver, typeParamVars}
    38  	argGet := make(map[*types.Var]func(*specexpr.Bindings) specexpr.Type)
    39  	ok := true
    40  	for _, v := range sFn.Params {
    41  		get := b.bindArg(v.Name(), v.Type())
    42  		ok = ok && (get != nil)
    43  		argGet[v] = get
    44  	}
    45  	for _, v := range sFn.Results {
    46  		get := b.bindArg(v.Name(), v.Type())
    47  		ok = ok && (get != nil)
    48  		argGet[v] = get
    49  	}
    50  	// Add requirements to the solver
    51  	for _, expr := range sFn.Requirements {
    52  		solver.Assert(expr)
    53  	}
    54  	if !ok {
    55  		return nil
    56  	}
    57  
    58  	// Find solutions.
    59  	defer func() {
    60  		p := recover()
    61  		if p != nil {
    62  			var buf strings.Builder
    63  			solver.Fprint(&buf)
    64  			panic(fmt.Sprintf("%s: %s\n%s", ctx.root.fset.Position(sFn.Pos), p, buf.String()))
    65  		}
    66  	}()
    67  	var funcs []*Func
    68  	for soln, err := range solver.Solve() {
    69  		if err != nil {
    70  			ctx.errorf("%s", err)
    71  			continue
    72  		}
    73  
    74  		fn := sFn.instantiate(ctx, soln, argGet)
    75  		if fn == nil {
    76  			continue
    77  		}
    78  		fn.typeParamVars = typeParamVars
    79  
    80  		funcs = append(funcs, fn)
    81  	}
    82  
    83  	if len(funcs) == 0 {
    84  		ctx.errorf("impossible constraints (try -f %s -trace)", sFn.Name)
    85  	}
    86  
    87  	return funcs
    88  }
    89  
    90  func (sFn *specFunc) instantiate(ctx context, b *specexpr.Bindings, argGet map[*types.Var]func(*specexpr.Bindings) specexpr.Type) *Func {
    91  	var f Func
    92  
    93  	f.Commutative = sFn.Commutative
    94  	f.Category = sFn.Category
    95  	f.Pos = sFn.Pkg.Fset.Position(sFn.Pos)
    96  	f.specFunc = sFn
    97  	f.instance = b
    98  
    99  	// Function or method?
   100  	var method bool
   101  	if len(sFn.Params) > 0 {
   102  		if t, ok := sFn.Params[0].Type().(*types.Named); ok {
   103  			if t.Origin() == sFn.Pkg.VecType {
   104  				method = true
   105  			}
   106  		}
   107  	}
   108  
   109  	// Instantiate name
   110  	name := sFn.NameTmpl.expand(func(s string) string {
   111  		val := b.Get(specexpr.Variable(s))
   112  		if val == nil {
   113  			ctx.errorf("unknown variable %q in function name", s)
   114  			return ""
   115  		}
   116  		str := fmt.Sprint(val)
   117  		// Make sure str starts with an upper-case letter so it maintains
   118  		// CamelCase in the overall identifier.
   119  		str = strings.ToTitle(str[:1]) + str[1:]
   120  		return str
   121  	})
   122  	f.Name = name
   123  
   124  	// Instantiate doc
   125  	doc := sFn.Doc.expand(func(s string) string {
   126  		val := b.Get(specexpr.Variable(s))
   127  		if val == nil {
   128  			ctx.errorf("unknown variable %q in doc", s)
   129  			return ""
   130  		}
   131  		return fmt.Sprint(val)
   132  	})
   133  	// Replace name in doc
   134  	if f.Name == sFn.Name {
   135  		f.Doc = doc
   136  	} else {
   137  		f.Doc = regexp.MustCompile(`\b`+regexp.QuoteMeta(sFn.Name)+`\b`).ReplaceAllLiteralString(doc, f.Name)
   138  	}
   139  
   140  	// Instantiate parameter and result types
   141  	//
   142  	// TODO: Should the loader keep these grouped like the original source so
   143  	// the transformed version keeps the same grouping (modulo pulling off the
   144  	// receiver)?
   145  	for _, v := range sFn.Params {
   146  		t := argGet[v](b)
   147  		f.In = append(f.In, Arg{v.Name(), t})
   148  	}
   149  	if method && len(f.In) > 0 {
   150  		f.Recv = f.In[0]
   151  		f.In = f.In[1:]
   152  	}
   153  	for _, v := range sFn.Results {
   154  		t := argGet[v](b)
   155  		f.Out = append(f.Out, Arg{v.Name(), t})
   156  	}
   157  
   158  	return &f
   159  }
   160
View as plain text