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.specFunc = sFn
    94  	f.instance = b
    95  
    96  	// Function or method?
    97  	var method bool
    98  	if len(sFn.Params) > 0 {
    99  		if t, ok := sFn.Params[0].Type().(*types.Named); ok {
   100  			if t.Origin() == sFn.Pkg.VecType {
   101  				method = true
   102  			}
   103  		}
   104  	}
   105  
   106  	// Instantiate name
   107  	name := sFn.NameTmpl.expand(func(s string) string {
   108  		val := b.Get(specexpr.Variable(s))
   109  		if val == nil {
   110  			ctx.errorf("unknown variable %q in function name", s)
   111  			return ""
   112  		}
   113  		str := fmt.Sprint(val)
   114  		// Make sure str starts with an upper-case letter so it maintains
   115  		// CamelCase in the overall identifier.
   116  		str = strings.ToTitle(str[:1]) + str[1:]
   117  		return str
   118  	})
   119  	f.Name = name
   120  
   121  	// Instantiate doc
   122  	doc := sFn.Doc.expand(func(s string) string {
   123  		val := b.Get(specexpr.Variable(s))
   124  		if val == nil {
   125  			ctx.errorf("unknown variable %q in doc", s)
   126  			return ""
   127  		}
   128  		return fmt.Sprint(val)
   129  	})
   130  	// Replace name in doc
   131  	if f.Name == sFn.Name {
   132  		f.Doc = doc
   133  	} else {
   134  		f.Doc = regexp.MustCompile(`\b`+regexp.QuoteMeta(sFn.Name)+`\b`).ReplaceAllLiteralString(doc, f.Name)
   135  	}
   136  
   137  	// Instantiate parameter and result types
   138  	//
   139  	// TODO: Should the loader keep these grouped like the original source so
   140  	// the transformed version keeps the same grouping (modulo pulling off the
   141  	// receiver)?
   142  	for _, v := range sFn.Params {
   143  		t := argGet[v](b)
   144  		f.In = append(f.In, Arg{v.Name(), t})
   145  	}
   146  	if method && len(f.In) > 0 {
   147  		f.Recv = f.In[0]
   148  		f.In = f.In[1:]
   149  	}
   150  	for _, v := range sFn.Results {
   151  		t := argGet[v](b)
   152  		f.Out = append(f.Out, Arg{v.Name(), t})
   153  	}
   154  
   155  	return &f
   156  }
   157
View as plain text