feat: add label in macro ite, add case in linker, add example test fo… (#1343)

* feat: add label in macro ite, add case in linker, add example test for ternary operator using a const

* feat: clean and add test

* feat: panic for register on rhs

* feat: factorize cases for linker

* feat: syntax error vs panic + fix linting

* feat: use ternary feature in TRM module

* fix: linker

Author: amkCha

pkg/asm/assembler/linker.go

@@ -203,14 +203,21 @@ func (p *Linker) linkInstruction(insn macro.Instruction, buses map[uint]io.Bus)
 		return p.linkExprs(insn.Sources)
 	case *macro.IfGoto:
 		if insn.Label != "" {
-			deats, ok := p.constmap[module.NewName(insn.Label, 1)]
-			//
-			if !ok {
-				msg := fmt.Sprintf("unknown register or constant \"%s\"", insn.Label)
-				return p.srcmap.SyntaxError(insn, msg)
+			value, errmsg := p.getLabelValue(insn.Label)
+			if errmsg != "" {
+				return p.srcmap.SyntaxError(insn, errmsg)
+			} else {
+				insn.Constant = value
+			}
+		}
+	case *macro.IfThenElse:
+		if insn.Label != "" {
+			value, errmsg := p.getLabelValue(insn.Label)
+			if errmsg != "" {
+				return p.srcmap.SyntaxError(insn, errmsg)
+			} else {
+				insn.Right = value
 			}
-			//
-			insn.Constant = deats.Left
 		}
 	default:
 		// continue
@@ -257,6 +264,21 @@ func (p *Linker) linkExprs(es []macro.Expr) *source.SyntaxError {
 	return nil
 }
 
+func (p *Linker) getLabelValue(l string) (big.Int, string) {
+	var deats util.Pair[big.Int, uint]
+	//
+	var ok bool
+	//
+	deats, ok = p.constmap[module.NewName(l, 1)]
+	//
+	if !ok {
+		msg := fmt.Sprintf("unknown register or constant \"%s\"", l)
+		return big.Int{}, msg
+	}
+	//
+	return deats.Left, ""
+}
+
 // Get the local bus declared for the given function, either by allocating a new
 // bus (if was not already allocated) or returning the existing bus (if it was
 // previously allocated).  Allocating a new bus requires allocating

pkg/asm/assembler/parser.go

@@ -573,7 +573,9 @@ func (p *Parser) parseTernaryRhs(targets []io.RegisterId, env *Environment) (mac
 	var (
 		errs        []source.SyntaxError
 		lhs         io.RegisterId
+		rhsExpr     macro.Expr
 		rhs, tb, fb big.Int
+		label       string
 		cond        uint8
 	)
 	// Parse left hand side
@@ -585,9 +587,19 @@ func (p *Parser) parseTernaryRhs(targets []io.RegisterId, env *Environment) (mac
 		return nil, errs
 	}
 	// Parse right hand side
-	if rhs, errs = p.parseNumber(env); len(errs) > 0 {
+	if rhsExpr, errs = p.parseAtomicExpr(env); len(errs) > 0 {
 		return nil, errs
 	}
+	// Dispatch on rhs expression form
+	switch e := rhsExpr.(type) {
+	case *expr.Const:
+		rhs = e.Constant
+		label = e.Label
+	case *expr.RegAccess:
+		// We can invoke (p.index - 1) as we are in the case of a ternary operator
+		// Checks are already performed to have a lhs
+		return nil, p.syntaxErrors(p.tokens[p.index-1], "ternary operator does not support register on the rhs")
+	}
 	// expect question mark
 	if _, errs = p.expect(QMARK); len(errs) > 0 {
 		return nil, errs
@@ -610,6 +622,7 @@ func (p *Parser) parseTernaryRhs(targets []io.RegisterId, env *Environment) (mac
 		Cond:    cond,
 		Left:    lhs,
 		Right:   rhs,
+		Label:   label,
 		Then:    tb,
 		Else:    fb,
 	}, nil

pkg/asm/io/macro/ite.go

@@ -31,6 +31,8 @@ type IfThenElse struct {
 	Left io.RegisterId
 	// Right-hand side
 	Right big.Int
+	// Constant label
+	Label string
 	// Then/Else branches
 	Then, Else big.Int
 }

pkg/test/assembly_util_test.go

@@ -62,3 +62,8 @@ func Test_AsmUtil_SetByte(t *testing.T) {
 func Test_AsmUtil_Signextend(t *testing.T) {
 	util.CheckWithFields(t, false, "asm/util/signextend", util.ASM_MAX_PADDING, field.BLS12_377, field.KOALABEAR_16)
 }
+
+// Note: field.KOALABEAR_16 is omitted as it's causing an issue with register splitting for the ternary operator
+func Test_AsmUtil_Ternary(t *testing.T) {
+	util.CheckWithFields(t, false, "asm/util/ternary", util.ASM_MAX_PADDING, field.BLS12_377)
+}

testdata/asm/bench/trm.zkasm

@@ -1,5 +1,5 @@
 const MAX_PRC_CONTIGUOUS_ADDRESS = 0x11
-;; const P256_VERIFY_ADDRESS = 0x100
+const P256_VERIFY_ADDRESS = 0x100
 
 ;; The address trimming module does two things: (i) reduce 256bit
 ;; addresses into u160 addresses (i.e. modulo 2^160); (b) identify
@@ -38,7 +38,7 @@ fn trm(RAW_ADDRESS u256) -> (ADDRESS_HI u32, IS_PRECOMPILE u1) {
   IS_PRECOMPILE=1
   return
 check_p256_verify:
-   IS_PRECOMPILE = low == 0x100 ? 1 : 0
+   IS_PRECOMPILE = low == P256_VERIFY_ADDRESS ? 1 : 0
    return
 exit_0:
   ;; no, not precompile

testdata/asm/util/ternary.accepts

@@ -0,0 +1,3 @@
+{"ternary_test": { "value": [199], "result": [0] }}
+{"ternary_test": { "value": [200], "result": [1] }}
+{"ternary_test": { "value": [201], "result": [0] }}

testdata/asm/util/ternary.zkasm

@@ -0,0 +1,7 @@
+const TWO_HUNDRED = 200
+
+;; ternary_test is using a value u128 to avoid the register splitting error for now
+fn ternary_test(value u128) -> (result u1) {
+    result = value == TWO_HUNDRED ? 1 : 0
+    return
+}