parse.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  	"strconv"
    10  	"strings"
    11  )
    12  
    13  // ParseExpr parses a constraint expression.
    14  func ParseExpr(x string) (Expr, error) {
    15  	var e Expr
    16  	var err error
    17  
    18  	func() {
    19  		defer func() {
    20  			v := recover()
    21  			if v2, ok := v.(*parseError); ok {
    22  				err = v2
    23  				return
    24  			} else if v != nil {
    25  				panic(v)
    26  			}
    27  		}()
    28  		p := &parser{s: x}
    29  		p.skipSpace()
    30  		e = p.parseExpr()
    31  		if p.pos < len(p.s) {
    32  			p.fail("unexpected trailing characters")
    33  		}
    34  	}()
    35  	return e, err
    36  }
    37  
    38  type parser struct {
    39  	s   string
    40  	pos int
    41  }
    42  
    43  type parseError struct {
    44  	msg  string
    45  	args []any
    46  	pos  int
    47  }
    48  
    49  func (e *parseError) Error() string {
    50  	return fmt.Sprintf("%s at %d", fmt.Sprintf(e.msg, e.args...), 1+e.pos)
    51  }
    52  
    53  func (p *parser) fail(msg string, args ...any) {
    54  	panic(&parseError{msg: msg, args: args, pos: p.pos})
    55  }
    56  
    57  func (p *parser) failAt(pos int, msg string, args ...any) {
    58  	panic(&parseError{msg: msg, args: args, pos: pos})
    59  }
    60  
    61  func (p *parser) skipSpace() {
    62  	for p.pos < len(p.s) && (p.s[p.pos] == ' ' || p.s[p.pos] == '\t' || p.s[p.pos] == '\r' || p.s[p.pos] == '\n') {
    63  		p.pos++
    64  	}
    65  }
    66  
    67  func (p *parser) peek() byte {
    68  	if p.pos >= len(p.s) {
    69  		return 0
    70  	}
    71  	return p.s[p.pos]
    72  }
    73  
    74  // consume consumes one or more characters matching pred. It does NOT consume
    75  // whitespace.
    76  func (p *parser) consume(pred func(c byte) bool) (string, bool) {
    77  	if p.pos >= len(p.s) || !pred(p.s[p.pos]) {
    78  		return "", false
    79  	}
    80  	start := p.pos
    81  	p.pos++
    82  	for p.pos < len(p.s) && pred(p.s[p.pos]) {
    83  		p.pos++
    84  	}
    85  	return p.s[start:p.pos], true
    86  }
    87  
    88  func isDigit(c byte) bool {
    89  	return c >= '0' && c <= '9'
    90  }
    91  
    92  func isAlpha(c byte) bool {
    93  	return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')
    94  }
    95  
    96  // try either consumes s followed zero or more whitespace and returns true, or
    97  // does nothing and returns false.
    98  func (p *parser) try(s string) bool {
    99  	if !strings.HasPrefix(p.s[p.pos:], s) {
   100  		return false
   101  	}
   102  	p.pos += len(s)
   103  	p.skipSpace()
   104  	return true
   105  }
   106  
   107  func (p *parser) peekShape() bool {
   108  	// Peek base type
   109  	if p.peek() == '{' {
   110  		return true
   111  	}
   112  	start := p.pos
   113  	if _, ok := p.consume(isAlpha); !ok {
   114  		return false
   115  	}
   116  	// Peek N
   117  	ok := p.pos < len(p.s) && (isDigit(p.s[p.pos]) || p.s[p.pos] == '{')
   118  	p.pos = start
   119  	return ok
   120  }
   121  
   122  func (p *parser) parseExpr() Expr {
   123  	return p.parseComparison()
   124  }
   125  
   126  func (p *parser) parseComparison() Expr {
   127  	x := p.parseMulDiv()
   128  	var op BinOp
   129  	switch {
   130  	case p.try("="):
   131  		op = OpEqual
   132  	case p.try("!="):
   133  		op = OpNotEqual
   134  	case p.try(">="):
   135  		op = OpGreaterOrEqual
   136  	case p.try(">"):
   137  		op = OpGreaterThan
   138  	case p.try("<="):
   139  		op = OpLessOrEqual
   140  	case p.try("<"):
   141  		op = OpLessThan
   142  	default:
   143  		return x
   144  	}
   145  
   146  	y := p.parseMulDiv()
   147  	return &BinExpr{op, x, y}
   148  }
   149  
   150  func (p *parser) parseMulDiv() Expr {
   151  	x := p.parsePrimary()
   152  loop:
   153  	for {
   154  		var op BinOp
   155  		switch {
   156  		case p.try("*"):
   157  			op = OpTimes
   158  		case p.try("/"):
   159  			op = OpDiv
   160  		default:
   161  			break loop
   162  		}
   163  		y := p.parsePrimary()
   164  		x = &BinExpr{op, x, y}
   165  	}
   166  	return x
   167  }
   168  
   169  func (p *parser) parsePrimary() Expr {
   170  	if p.try("(") {
   171  		e := p.parseExpr()
   172  		if !p.try(")") {
   173  			p.fail("expected ')'")
   174  		}
   175  		return e
   176  	}
   177  
   178  	// Shape
   179  	if p.peekShape() {
   180  		return p.parseSymShape()
   181  	}
   182  
   183  	// Number literal
   184  	b := p.peek()
   185  	if isDigit(b) {
   186  		return p.parseNumber()
   187  	}
   188  
   189  	// Variable
   190  	if isAlpha(b) {
   191  		name, _ := p.consume(isAlpha)
   192  		p.skipSpace()
   193  		return Variable(name)
   194  	}
   195  
   196  	if b == 0 {
   197  		p.fail("unexpected end")
   198  	}
   199  	p.fail("unexpected character '%c'", b)
   200  	panic("not reachable")
   201  }
   202  
   203  func (p *parser) parseNumber() Int {
   204  	nStr, ok := p.consume(isDigit)
   205  	if !ok {
   206  		p.fail("expected number")
   207  	}
   208  	num, err := strconv.Atoi(nStr)
   209  	if err != nil {
   210  		p.fail("%s", err)
   211  	}
   212  	p.skipSpace()
   213  	return Int(num)
   214  }
   215  
   216  // - BaseNxL: A fixed vector with L lanes. E.g., Int32x4
   217  // - BaseNs: A scalable vector. E.g., Float32s
   218  // - BaseNwW: A fixed vector of width W. E.g., Int32w128 (same as Int32x4)
   219  // - MaskNxL, MaskNs, or MaskNwW: Similar, but describes a mask.
   220  // - baseN: A scalar. E.g., uint8
   221  func (p *parser) parseSymShape() *Apply {
   222  	var b, n Expr
   223  
   224  	trySymPart := func() Expr {
   225  		openPos := p.pos
   226  		if !p.try("{") {
   227  			return nil
   228  		}
   229  
   230  		x := p.parseExpr()
   231  		if !p.try("}") {
   232  			p.failAt(openPos, "'{' missing close '}' in symbolic shape")
   233  		}
   234  		return x
   235  	}
   236  
   237  	// Base
   238  	if b = trySymPart(); b == nil {
   239  		base, ok := p.consume(isAlpha)
   240  		if !ok {
   241  			p.fail("expected shape base name matching [a-zA-Z]+")
   242  		}
   243  		b = &Literal{strings.ToLower(base)}
   244  	}
   245  
   246  	// Element size
   247  	if n = trySymPart(); n == nil {
   248  		n = p.parseNumber()
   249  	}
   250  
   251  	elem := MakeBasic(b, n)
   252  
   253  	// Width
   254  	var x *Apply
   255  	// Don't use p.try here because that will skip whitespace.
   256  	switch p.peek() {
   257  	case 's':
   258  		p.pos++
   259  		x = MakeVector(elem, VW())
   260  	case 'x':
   261  		p.pos++
   262  		l := trySymPart()
   263  		if l == nil {
   264  			// TODO: Disallow width-rounding in this case?
   265  			l = p.parseNumber()
   266  		}
   267  		x = makeVectorL(elem, l)
   268  	case 'w':
   269  		p.pos++
   270  		w := trySymPart()
   271  		if w == nil {
   272  			w = p.parseNumber()
   273  		}
   274  		x = MakeVector(elem, w)
   275  	default:
   276  		// Scalar
   277  		x = elem
   278  	}
   279  
   280  	p.skipSpace()
   281  	return x
   282  }
   283
View as plain text