expr.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 specexpr
     6  
     7  import (
     8  	"fmt"
     9  	"iter"
    10  	"reflect"
    11  	"strings"
    12  	"sync"
    13  )
    14  
    15  type Expr interface {
    16  	String() string
    17  	eval(b *Bindings) (any, error)
    18  	preorder(yield func(Expr) bool) bool
    19  }
    20  
    21  // exprVars yields all Variable nodes in an Expr.
    22  func exprVars(e Expr) iter.Seq[Variable] {
    23  	return func(yield func(Variable) bool) {
    24  		e.preorder(func(e Expr) bool {
    25  			if v, ok := e.(Variable); ok {
    26  				return yield(v)
    27  			}
    28  			return true
    29  		})
    30  	}
    31  }
    32  
    33  // Literal is an [Expr] that evaluates to a literal value.
    34  //
    35  // For [Int] and [SymbolicWidth], you probably just want to use those types
    36  // directly. They're literal values, so you could wrap them in a Literal, but
    37  // they are valid expressions on their own.
    38  type Literal struct {
    39  	Val any
    40  }
    41  
    42  func (e *Literal) String() string {
    43  	return fmt.Sprint(e.Val)
    44  }
    45  func (e *Literal) eval(b *Bindings) (any, error) {
    46  	if i, ok := e.Val.(int); ok {
    47  		// The evaluator works with Nums, not ints directly.
    48  		return Int(i), nil
    49  	}
    50  	return e.Val, nil
    51  }
    52  func (e *Literal) preorder(yield func(Expr) bool) bool {
    53  	return yield(e)
    54  }
    55  
    56  // Variable is an [Expr] that evaluates to the value of the named variable.
    57  type Variable string
    58  
    59  func (e Variable) String() string {
    60  	return string(e)
    61  }
    62  func (e Variable) eval(b *Bindings) (any, error) {
    63  	val := b.Get(e)
    64  	if val == nil {
    65  		panic(fmt.Errorf("variable %s not solved", e))
    66  	}
    67  	return val, nil
    68  }
    69  func (e Variable) preorder(yield func(Expr) bool) bool {
    70  	return yield(e)
    71  }
    72  
    73  // Func is a function that can be used in an expression. Use the [Func.Apply]
    74  // method to create an [Expr].
    75  type Func struct {
    76  	Name string
    77  	Func func([]any) (any, error)
    78  }
    79  
    80  func MakeFunc1[T any](name string, fn func(T) (any, error)) func(e Expr) *Apply {
    81  	f := &Func{
    82  		Name: name,
    83  		Func: func(a []any) (any, error) {
    84  			if len(a) != 1 {
    85  				panic(fmt.Sprintf("%s: got %d arguments, want %d", name, len(a), 1))
    86  			}
    87  			v, ok := a[0].(T)
    88  			if !ok {
    89  				panic(fmt.Sprintf("%s: argument is %T, want %T", name, a[0], *new(T)))
    90  			}
    91  			return fn(v)
    92  		},
    93  	}
    94  	return func(e Expr) *Apply {
    95  		return f.Apply(e)
    96  	}
    97  }
    98  func MakeFunc2[T, U any](name string, fn func(T, U) (any, error)) func(e1, e2 Expr) *Apply {
    99  	f := &Func{
   100  		Name: name,
   101  		Func: func(a []any) (any, error) {
   102  			if len(a) != 2 {
   103  				panic(fmt.Sprintf("%s: got %d arguments, want %d", name, len(a), 2))
   104  			}
   105  			v1, ok := a[0].(T)
   106  			if !ok {
   107  				panic(fmt.Sprintf("%s: argument is %T, want %T", name, a[0], *new(T)))
   108  			}
   109  			v2, ok := a[1].(U)
   110  			if !ok {
   111  				panic(fmt.Sprintf("%s: argument is %T, want %T", name, a[1], *new(U)))
   112  			}
   113  			return fn(v1, v2)
   114  		},
   115  	}
   116  	return func(e1, e2 Expr) *Apply {
   117  		return f.Apply(e1, e2)
   118  	}
   119  }
   120  
   121  func MakeFunc(name string, fn func([]any) (any, error)) *Func {
   122  	return &Func{name, fn}
   123  }
   124  
   125  type fieldGetterKey struct {
   126  	rt        reflect.Type
   127  	fieldName string
   128  }
   129  
   130  var fieldGetters sync.Map
   131  
   132  // MakeField returns a Func that projects field fieldName from a value of type
   133  // T. The returned functions are memoized, so only one *Func is created per type
   134  // and field and this is efficient to call repeatedly.
   135  func MakeField[T any](fieldName string) *Func {
   136  	rt := reflect.TypeFor[T]()
   137  	key := fieldGetterKey{rt, fieldName}
   138  	get, ok := fieldGetters.Load(key)
   139  	if !ok {
   140  		f, ok := rt.FieldByName(fieldName)
   141  		if !ok {
   142  			panic(fmt.Sprintf("no such field %s in type %s", fieldName, rt))
   143  		}
   144  		getF := MakeFunc(rt.Name()+"."+fieldName, func(a []any) (any, error) {
   145  			if len(a) != 1 {
   146  				panic("expected exactly 1 argument")
   147  			}
   148  			var rv reflect.Value
   149  			if _, ok := a[0].(T); ok {
   150  				rv = reflect.ValueOf(a[0])
   151  			} else if _, ok := a[0].(*T); ok {
   152  				rv = reflect.ValueOf(a[0]).Elem()
   153  			} else {
   154  				panic(fmt.Sprintf("argument is %T, want %s", a[0], rt))
   155  			}
   156  			return rv.FieldByIndex(f.Index).Interface(), nil
   157  		})
   158  		get, _ = fieldGetters.LoadOrStore(key, getF)
   159  	}
   160  	return get.(*Func)
   161  }
   162  
   163  func (f *Func) Apply(args ...Expr) *Apply {
   164  	return &Apply{f, args}
   165  }
   166  
   167  // Apply is an [Expr] that applies a [Func] to a sequence of arguments.
   168  type Apply struct {
   169  	Func *Func
   170  	Args []Expr
   171  }
   172  
   173  func (e *Apply) String() string {
   174  	var buf strings.Builder
   175  	buf.WriteString(e.Func.Name)
   176  	buf.WriteByte('(')
   177  	for i, x := range e.Args {
   178  		if i > 0 {
   179  			buf.WriteString(", ")
   180  		}
   181  		buf.WriteString(x.String())
   182  	}
   183  	buf.WriteByte(')')
   184  	return buf.String()
   185  }
   186  func (e *Apply) eval(b *Bindings) (any, error) {
   187  	vals := make([]any, 0, 16)
   188  	for _, arg := range e.Args {
   189  		val, err := arg.eval(b)
   190  		if err != nil {
   191  			return nil, err
   192  		}
   193  		vals = append(vals, val)
   194  	}
   195  	return e.Func.Func(vals)
   196  }
   197  func (e *Apply) preorder(yield func(Expr) bool) bool {
   198  	if !yield(e) {
   199  		return false
   200  	}
   201  	for _, a := range e.Args {
   202  		if !a.preorder(yield) {
   203  			return false
   204  		}
   205  	}
   206  	return true
   207  }
   208  
   209  // BinExpr is a binary [Expr].
   210  type BinExpr struct {
   211  	Op   BinOp
   212  	X, Y Expr
   213  }
   214  
   215  func (e *BinExpr) String() string {
   216  	op := "???"
   217  	if int(e.Op) < len(opStrings) && opStrings[e.Op] != "" {
   218  		op = opStrings[e.Op]
   219  	}
   220  	return e.X.String() + op + e.Y.String()
   221  }
   222  func (e *BinExpr) preorder(yield func(Expr) bool) bool {
   223  	return yield(e) && e.X.preorder(yield) && e.Y.preorder(yield)
   224  }
   225  
   226  type BinOp byte
   227  
   228  const (
   229  	_       BinOp = iota
   230  	OpTimes       // int = int * int or Width = int * Width (or Width * int)
   231  	OpDiv         // int = int / int (must be exact) or Width = Width / int
   232  
   233  	// Comparison operators
   234  	OpEqual          // bool = expr = expr
   235  	OpNotEqual       // bool = expr != expr
   236  	OpGreaterThan    // bool = expr > expr
   237  	OpLessThan       // bool = expr < expr
   238  	OpGreaterOrEqual // bool = expr >= expr
   239  	OpLessOrEqual    // bool = expr <= expr
   240  )
   241  
   242  var opStrings = [...]string{
   243  	OpTimes: "*",
   244  	OpDiv:   "/",
   245  
   246  	OpEqual:          "=",
   247  	OpNotEqual:       "!=",
   248  	OpGreaterThan:    ">",
   249  	OpLessThan:       "<",
   250  	OpGreaterOrEqual: ">=",
   251  	OpLessOrEqual:    "<=",
   252  }
   253  
   254  func (e *BinExpr) eval(b *Bindings) (any, error) {
   255  	xVal, err := e.X.eval(b)
   256  	if err != nil {
   257  		return nil, err
   258  	}
   259  	yVal, err := e.Y.eval(b)
   260  	if err != nil {
   261  		return nil, err
   262  	}
   263  
   264  	xn, okX := xVal.(Num)
   265  	yn, okY := yVal.(Num)
   266  
   267  	switch e.Op {
   268  	case OpEqual, OpNotEqual:
   269  		if okX && okY {
   270  			// Fall through to numeric operations
   271  			break
   272  		}
   273  		// Otherwise, general equality
   274  		if reflect.TypeOf(xVal) != reflect.TypeOf(yVal) {
   275  			panic(fmt.Errorf("incompatible types for comparison: %T and %T", xVal, yVal))
   276  		}
   277  		if e.Op == OpEqual {
   278  			return xVal == yVal, nil
   279  		} else {
   280  			return xVal != yVal, nil
   281  		}
   282  	}
   283  
   284  	if !okX {
   285  		panic(fmt.Errorf("invalid type for %v: %v (%T)", e.Op, xVal, xVal))
   286  	}
   287  	if !okY {
   288  		panic(fmt.Errorf("invalid type for %v: %v (%T)", e.Op, yVal, yVal))
   289  	}
   290  
   291  	switch e.Op {
   292  	case OpTimes:
   293  		return xn.Mul(yn)
   294  
   295  	case OpDiv:
   296  		return xn.Div(yn)
   297  
   298  	case OpEqual, OpNotEqual, OpGreaterThan, OpLessThan, OpGreaterOrEqual, OpLessOrEqual:
   299  		return e.evalComparison(xn, yn)
   300  	}
   301  
   302  	panic("bad binop")
   303  }
   304  
   305  func (e *BinExpr) evalComparison(x, y Num) (any, error) {
   306  	res, ok := x.Compare(y)
   307  	if !ok {
   308  		// Incomparable
   309  		return e.Op == OpNotEqual, nil
   310  	}
   311  	switch e.Op {
   312  	case OpEqual:
   313  		return res == 0, nil
   314  	case OpNotEqual:
   315  		return res != 0, nil
   316  	case OpGreaterThan:
   317  		return res > 0, nil
   318  	case OpLessThan:
   319  		return res < 0, nil
   320  	case OpGreaterOrEqual:
   321  		return res >= 0, nil
   322  	case OpLessOrEqual:
   323  		return res <= 0, nil
   324  	}
   325  	panic("bad comparison operator")
   326  }
   327
View as plain text