From 2c5c3bb421a4bfa9dbda00b7c10c18d5cebadef1 Mon Sep 17 00:00:00 2001
From: Felix Lange <fjl@twurst.com>
Date: Wed, 27 Sep 2023 23:47:59 +0200
Subject: [PATCH] asm: add builtins .bitlen and .bytelen

---
 asm/builtins.go  | 25 +++++++++++++++++++++++++
 asm/eval_test.go |  7 +++++++
 2 files changed, 32 insertions(+)

diff --git a/asm/builtins.go b/asm/builtins.go
index 238798f..2ff7c5e 100644
--- a/asm/builtins.go
+++ b/asm/builtins.go
@@ -29,6 +29,8 @@ import (
 )
 
 var builtinMacros = map[string]builtinMacroFn{
+	"bitlen":    bitlenMacro,
+	"bytelen":   bytelenMacro,
 	"abs":       absMacro,
 	"address":   addressMacro,
 	"selector":  selectorMacro,
@@ -38,6 +40,29 @@ var builtinMacros = map[string]builtinMacroFn{
 
 type builtinMacroFn func(*evaluator, *evalEnvironment, *macroCallExpr) (*big.Int, error)
 
+func bitlenMacro(e *evaluator, env *evalEnvironment, call *macroCallExpr) (*big.Int, error) {
+	if err := call.checkArgCount(1); err != nil {
+		return nil, err
+	}
+	v, err := call.args[0].eval(e, env)
+	if err != nil {
+		return nil, err
+	}
+	return big.NewInt(int64(v.BitLen())), nil
+}
+
+func bytelenMacro(e *evaluator, env *evalEnvironment, call *macroCallExpr) (*big.Int, error) {
+	if err := call.checkArgCount(1); err != nil {
+		return nil, err
+	}
+	v, err := call.args[0].eval(e, env)
+	if err != nil {
+		return nil, err
+	}
+	bytes := (v.BitLen() + 7) / 8
+	return big.NewInt(int64(bytes)), nil
+}
+
 func absMacro(e *evaluator, env *evalEnvironment, call *macroCallExpr) (*big.Int, error) {
 	if err := call.checkArgCount(1); err != nil {
 		return nil, err
diff --git a/asm/eval_test.go b/asm/eval_test.go
index f1f07ab..05c360c 100644
--- a/asm/eval_test.go
+++ b/asm/eval_test.go
@@ -58,6 +58,13 @@ var evalTests = []evalTest{
 	{expr: `"A"`, result: "65"},
 	{expr: `"foo"`, result: "6713199"},
 	// builtins
+	{expr: `.bitlen(0)`, result: "0"},
+	{expr: `.bitlen(0xff)`, result: "8"},
+	{expr: `.bitlen(0x01ff)`, result: "9"},
+	{expr: `.bytelen(0)`, result: "0"},
+	{expr: `.bytelen(0xff)`, result: "1"},
+	{expr: `.bytelen(0x01ff)`, result: "2"},
+	{expr: `.bytelen("foobar")`, result: "6"},
 	{expr: `.abs(0 - 10)`, result: "10"},
 	{expr: `.sha256("text")`, result: "68832153269555879243704685382415794081420120252170153643880971663484982053329"},
 	{expr: `.sha256(33)`, result: "84783983549258160669137366770885509408211009960610860350324922232842582506338"},