types.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  	"strconv"
    13  	"strings"
    14  	"sync"
    15  
    16  	"golang.org/x/tools/go/types/typeutil"
    17  )
    18  
    19  // constraintToDomain enumerates all types explicitly listed as satisfying
    20  // constraint (which must be a type parameter constraint), and translates them
    21  // to a [specexpr] domain.
    22  func constraintToDomain(pkg *specPackage, constraint types.Type) ([]any, error) {
    23  	var vals []any
    24  	for _, typ := range typeSet(constraint) {
    25  		elem, ok := pkg.TypeElems[typ]
    26  		if ok {
    27  			vals = append(vals, elem)
    28  		} else if width, ok := pkg.TypeWidths[typ]; ok {
    29  			vals = append(vals, width)
    30  		} else {
    31  			return nil, fmt.Errorf("type %s satisfies constraint %s, but isn't a known shape type", typ, constraint)
    32  		}
    33  	}
    34  	return vals, nil
    35  }
    36  
    37  var typeSetMemo sync.Map
    38  
    39  // typeSet enumerates all concrete types in t's type set.
    40  func typeSet(t types.Type) []types.Type {
    41  	// In general types.Types are not comparable, but every type we're dealing
    42  	// with here has pointer identity, and also there's no correctness issue if
    43  	// we miss in the memo.
    44  	ts, ok := typeSetMemo.Load(t)
    45  	if !ok {
    46  		var o orderedTypeSet
    47  		typeSet1(t, &o)
    48  		ts, _ = typeSetMemo.LoadOrStore(t, o.types)
    49  	}
    50  	return ts.([]types.Type)
    51  }
    52  
    53  type orderedTypeSet struct {
    54  	types []types.Type
    55  	set   typeutil.Map
    56  }
    57  
    58  func (set *orderedTypeSet) add(t types.Type) {
    59  	if set.set.At(t) == nil {
    60  		set.set.Set(t, true)
    61  		set.types = append(set.types, t)
    62  	}
    63  }
    64  
    65  func (set *orderedTypeSet) intersect(o orderedTypeSet) {
    66  	i, j := 0, 0
    67  	for ; i < len(set.types); i++ {
    68  		t := set.types[i]
    69  		if o.set.At(t) != nil {
    70  			set.types[j] = t
    71  			j++
    72  		} else {
    73  			set.set.Delete(t)
    74  		}
    75  	}
    76  	set.types = set.types[:j]
    77  }
    78  
    79  func typeSet1(t types.Type, o *orderedTypeSet) {
    80  	switch u := t.Underlying().(type) {
    81  	case *types.Interface:
    82  		switch u.NumEmbeddeds() {
    83  		case 0:
    84  			return
    85  		case 1:
    86  			// Fast path for common case
    87  			typeSet1(u.EmbeddedType(0), o)
    88  			return
    89  		}
    90  		var intersection orderedTypeSet
    91  		first := true
    92  		for etyp := range u.EmbeddedTypes() {
    93  			if first {
    94  				typeSet1(etyp, &intersection)
    95  				first = false
    96  			} else {
    97  				var tmp orderedTypeSet
    98  				typeSet1(etyp, &tmp)
    99  				intersection.intersect(tmp)
   100  			}
   101  		}
   102  		// TODO: This doesn't check method satisfaction. We could do that with
   103  		// types.Satisfies filter here, but it doesn't matter for our needs.
   104  		for _, etyp := range intersection.types {
   105  			o.add(etyp)
   106  		}
   107  
   108  	case *types.Union:
   109  		for term := range u.Terms() {
   110  			typeSet1(term.Type(), o)
   111  		}
   112  
   113  	default:
   114  		o.add(t)
   115  	}
   116  }
   117  
   118  var basicRe = regexp.MustCompile(`^([a-z]+|Mask)([0-9]*)$`)
   119  
   120  // shapeElemType parses a basic or mask element spec type.
   121  func shapeElemType(t types.Type) specexpr.Basic {
   122  	var name string
   123  	switch t := t.(type) {
   124  	case *types.Basic: // E.g., uint32
   125  		name = t.Name()
   126  	case *types.Named: // E.g., Mask16
   127  		name = t.Obj().Name()
   128  	default:
   129  		panic(fmt.Sprintf("not a shape element type: %s", t))
   130  	}
   131  	m := basicRe.FindStringSubmatch(name)
   132  	if m == nil {
   133  		panic(fmt.Sprintf("failed to parse element type %s", name))
   134  	}
   135  	bits := 0
   136  	if m[2] != "" {
   137  		bits, _ = strconv.Atoi(m[2])
   138  	}
   139  	return specexpr.Basic{Base: m[1], Bits: specexpr.Int(bits)}
   140  }
   141  
   142  // shapeWidthVal parses a spec width type.
   143  func shapeWidthVal(t types.Type) specexpr.Num {
   144  	named, ok := t.(*types.Named)
   145  	if !ok {
   146  		panic(fmt.Sprintf("not a shape width type: %s", t))
   147  	}
   148  	name := named.Obj().Name()
   149  	if name == "WidthScalable" {
   150  		return specexpr.VW()
   151  	}
   152  	var err error
   153  	if suffix, ok := strings.CutPrefix(name, "Width"); ok {
   154  		var val int
   155  		val, err = strconv.Atoi(suffix)
   156  		if err == nil {
   157  			return specexpr.Int(val)
   158  		}
   159  	} else {
   160  		err = fmt.Errorf("does not start with 'Width'")
   161  	}
   162  	panic(fmt.Sprintf("parsing width type %s: %s", t, err))
   163  }
   164  
   165  var basicToBasic = map[specexpr.Basic]types.BasicKind{
   166  	{Base: "int", Bits: 0}:    types.Int,
   167  	{Base: "int", Bits: 8}:    types.Int8,
   168  	{Base: "int", Bits: 16}:   types.Int16,
   169  	{Base: "int", Bits: 32}:   types.Int32,
   170  	{Base: "int", Bits: 64}:   types.Int64,
   171  	{Base: "uint", Bits: 0}:   types.Uint,
   172  	{Base: "uint", Bits: 8}:   types.Uint8,
   173  	{Base: "uint", Bits: 16}:  types.Uint16,
   174  	{Base: "uint", Bits: 32}:  types.Uint32,
   175  	{Base: "uint", Bits: 64}:  types.Uint64,
   176  	{Base: "float", Bits: 32}: types.Float32,
   177  	{Base: "float", Bits: 64}: types.Float64,
   178  }
   179  
   180  func specTypeToType(pkg *specPackage, t specexpr.Type) types.Type {
   181  	switch t := t.(type) {
   182  	case specexpr.Basic:
   183  		var t2 types.Type
   184  		if t.Base == "Mask" {
   185  			t2 = pkg.ElemTypes[t]
   186  		} else {
   187  			if kind, ok := basicToBasic[t]; ok {
   188  				t2 = types.Typ[kind]
   189  			}
   190  		}
   191  		if t2 == nil {
   192  			panic(fmt.Sprintf("unknown basic type %s", t))
   193  		}
   194  		return t2
   195  	}
   196  	// TODO: Implement other specexpr.Type types if we need them
   197  	panic(fmt.Sprintf("unimplemented specTypeToType for %T", t))
   198  }
   199  
   200  type argBinder struct {
   201  	ctx        context
   202  	pkg        *specPackage
   203  	s          *specexpr.Solver
   204  	typeParams map[*types.TypeParam]specexpr.Variable
   205  }
   206  
   207  // bindArg assigns all solver variables related to an argument called "name" of
   208  // type t. It returns a function that retrieves the resolved type, or nil on
   209  // error.
   210  func (b *argBinder) bindArg(name string, t types.Type) func(*specexpr.Bindings) specexpr.Type {
   211  	expr := b.bind1(name, t)
   212  	if expr == nil {
   213  		return nil
   214  	}
   215  	v := b.s.Assign(specexpr.Variable(name), expr)
   216  	return func(b *specexpr.Bindings) specexpr.Type {
   217  		return b.Get(v).(specexpr.Type)
   218  	}
   219  }
   220  
   221  // bind1 deconstructs t and binds any components to variables derived from
   222  // "name", and returns the (not yet bound!) expression for t. The caller is
   223  // expected to bind "name" to the returned expression, or pass it up. It works
   224  // this way so we can unwrap things like pointer and slice types without
   225  // creating intermediate names for each level.
   226  func (b *argBinder) bind1(name string, t types.Type) specexpr.Expr {
   227  	switch t := t.(type) {
   228  	case *types.Basic:
   229  		return &specexpr.Literal{Val: shapeElemType(t)}
   230  
   231  	case *types.Pointer:
   232  		elem := b.bind1(name, t.Elem())
   233  		if elem == nil {
   234  			return nil
   235  		}
   236  		return specexpr.MakePointer(elem)
   237  
   238  	case *types.Slice:
   239  		elem := b.bind1(name, t.Elem())
   240  		if elem == nil {
   241  			return nil
   242  		}
   243  		return specexpr.MakeSlice(elem)
   244  
   245  	case *types.Array:
   246  		elem := b.bind1(name, t.Elem())
   247  		if elem == nil {
   248  			return nil
   249  		}
   250  		return specexpr.MakeArray(elem, specexpr.Int(t.Len()))
   251  
   252  	case *types.Named:
   253  		if types.Identical(t.Origin(), b.pkg.VecType) {
   254  			xE, xW, _ := b.bindVecLike(name, t)
   255  			if xE == nil {
   256  				return nil
   257  			}
   258  			return specexpr.MakeVector(xE, xW)
   259  		}
   260  		if types.Identical(t.Origin(), b.pkg.ArrayType) {
   261  			xE, xW, xL := b.bindVecLike(name, t)
   262  			if xE == nil {
   263  				return nil
   264  			}
   265  			b.s.Assert(funcAssertFixed(xW))
   266  			return specexpr.MakeArray(xE, xL)
   267  		}
   268  		if types.Identical(t.Origin(), b.pkg.UintNType) {
   269  			xN := specexpr.Variable(name + "N")
   270  			b.s.Declare(xN, []any{8, 16, 32, 64})
   271  			b.s.Assert(funcAssertFixed(xN))
   272  			return specexpr.MakeBasic(&specexpr.Literal{Val: "uint"}, xN)
   273  		}
   274  
   275  	case *types.TypeParam:
   276  		return b.typeParams[t]
   277  	}
   278  
   279  	b.ctx.errorf("cannot convert spec type %s into API type", t)
   280  	return nil
   281  }
   282  
   283  var funcAssertFixed = specexpr.MakeFunc1("assertFixed", func(w specexpr.Num) (any, error) {
   284  	_, ok := w.(specexpr.Int)
   285  	return ok, nil
   286  })
   287  
   288  func (b *argBinder) bindVecLike(name string, t *types.Named) (xE, xW, xL specexpr.Expr) {
   289  	args := t.TypeArgs()
   290  	if args.Len() != 2 {
   291  		b.ctx.errorf("expected exactly 2 type arguments, got %d", args.Len())
   292  		return nil, nil, nil
   293  	}
   294  
   295  	// Assign the element type
   296  	elem := b.bind1(name+"E", args.At(0))
   297  	if elem == nil {
   298  		return nil, nil, nil
   299  	}
   300  	xE = b.s.Assign(specexpr.Variable(name+"E"), elem)
   301  
   302  	// Get the width
   303  	var wExpr specexpr.Expr
   304  	switch wt := args.At(1).(type) {
   305  	case *types.TypeParam:
   306  		wExpr = b.typeParams[wt]
   307  	case *types.Named:
   308  		wExpr = b.pkg.TypeWidths[wt]
   309  		if wExpr == nil {
   310  			b.ctx.errorf("width type argument is not a width")
   311  			return nil, nil, nil
   312  		}
   313  	default:
   314  		b.ctx.errorf("width type arguments not a type parameter or named type")
   315  		return nil, nil, nil
   316  	}
   317  	xW = b.s.Assign(specexpr.Variable(name+"W"), wExpr)
   318  
   319  	// Bind other variables
   320  	basicBase := specexpr.MakeField[specexpr.Basic]("Base")
   321  	basicBits := specexpr.MakeField[specexpr.Basic]("Bits")
   322  	b.s.Assign(specexpr.Variable(name+"B"), basicBase.Apply(xE))
   323  	xN := b.s.Assign(specexpr.Variable(name+"N"), basicBits.Apply(xE))
   324  	xL = b.s.Assign(specexpr.Variable(name+"L"), &specexpr.BinExpr{
   325  		Op: specexpr.OpDiv, X: xW, Y: xN,
   326  	})
   327  
   328  	return xE, xW, xL
   329  }
   330
View as plain text