feat: script builder

builder/builder.go

@@ -0,0 +1,152 @@
+// Builder to format numscript document
+package builder
+
+import (
+	"fmt"
+	"math/big"
+	"strings"
+)
+
+const indentStr = "  "
+
+func (p *pool[T]) getFreshId() int {
+	id := p.nextId
+	p.nextId += 1
+	return id
+}
+
+func (p *pool[T]) getItemId(elem T) int {
+	previousLookup, isElemInPool := p.elems[elem]
+	if !isElemInPool {
+		elemId := p.getFreshId()
+		p.elems[elem] = elemId
+		previousLookup = elemId
+	}
+	return previousLookup
+}
+
+type env struct {
+	builder      strings.Builder
+	accountsPool pool[string]
+	assetsPool   pool[string]
+	stringsPool  pool[string]
+	numbersPool  pool[*big.Int]
+	varsEnv      VarsEnv
+}
+
+func writeIndentation(env *env, w int) {
+	if w == 0 {
+		return
+	}
+
+	env.builder.Grow(w * len(indentStr))
+	for range w {
+		env.builder.WriteString(indentStr)
+	}
+}
+
+func newEnv() env {
+	return env{
+		accountsPool: newPool[string](),
+		assetsPool:   newPool[string](),
+		stringsPool:  newPool[string](),
+		numbersPool:  newPool[*big.Int](),
+		varsEnv:      VarsEnv{bindings: map[anyVar]string{}},
+	}
+}
+
+// The underlying type of any a pretty printing document
+type render = func(
+	env *env,
+
+	// The current width
+	w int,
+)
+
+func itemIdToName(id int, prefix string) string {
+	return fmt.Sprintf("%s_%d", prefix, id)
+}
+func accountToName(id int) string {
+	return itemIdToName(id, "account")
+}
+func assetToName(id int) string {
+	return itemIdToName(id, "asset")
+}
+func stringToName(id int) string {
+	return itemIdToName(id, "string")
+}
+func numberToName(id int) string {
+	return itemIdToName(id, "number")
+}
+
+type varRenderState struct {
+	hasVars       bool
+	sb            strings.Builder
+	knownBindings map[string]string
+}
+
+func renderVar[T comparable](
+	st *varRenderState,
+
+	typ string,
+	pool pool[T],
+	getVarName func(id int) string,
+	stringifyValue func(value T) string,
+) {
+	for key, id := range pool.elems {
+		varName := getVarName(id)
+		st.knownBindings[varName] = stringifyValue(key)
+	}
+
+	for id := range pool.nextId {
+		st.hasVars = true
+
+		varName := getVarName(id)
+
+		st.sb.WriteString(indentStr)
+		st.sb.WriteString(typ)
+		st.sb.WriteString(" $")
+		st.sb.WriteString(varName)
+		st.sb.WriteByte('\n')
+	}
+
+}
+
+func stringId(x string) string { return x }
+
+func renderVars(
+	st *varRenderState,
+	env *env,
+) string {
+
+	st.sb.WriteString("vars {\n")
+	renderVar(st, "account", env.accountsPool, accountToName, stringId)
+	renderVar(st, "string", env.stringsPool, stringToName, stringId)
+	renderVar(st, "asset", env.assetsPool, assetToName, stringId)
+	renderVar(st, "number", env.numbersPool, numberToName, func(bi *big.Int) string {
+		return bi.String()
+	})
+	st.sb.WriteString("}\n\n")
+
+	if !st.hasVars {
+		return ""
+	}
+
+	return st.sb.String()
+}
+
+// TODO double check this one (do we need to handle vars?)
+func BuildProgram(statements ...Statement) (map[string]string, VarsEnv, string) {
+	env := newEnv()
+	for _, stmt := range statements {
+		stmt(&env, 0)
+	}
+
+	st := varRenderState{
+		knownBindings: make(map[string]string),
+	}
+	// AFTER we've rendered the whole program, we can render the vars block
+	vars := renderVars(&st, &env)
+
+	return st.knownBindings, env.varsEnv, vars + env.builder.String()
+}

builder/builder_test.go

@@ -0,0 +1,244 @@
+package builder_test
+
+import (
+	"math/big"
+	"testing"
+
+	"github.com/formancehq/numscript/builder"
+	"github.com/gkampitakis/go-snaps/snaps"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSimpleSend(t *testing.T) {
+	stmt := builder.StmtSend(
+		builder.ExprMonetary(
+			builder.ExprAsset("USD/2"),
+			builder.ExprNumberBigInt(big.NewInt(42)),
+		),
+		builder.SrcAccount(
+			builder.ExprAccount("src"),
+		),
+		builder.DestAccount(
+			builder.ExprAccount("dest"),
+		),
+	)
+
+	_, _, script := builder.BuildProgram(stmt)
+	snaps.MatchInlineSnapshot(t, script, snaps.Inline(`vars {
+  account $account_0
+  account $account_1
+  asset $asset_0
+}
+
+send [$asset_0 42] (
+  source = $account_0
+  destination = $account_1
+)`))
+}
+
+func TestInorder(t *testing.T) {
+	stmt := builder.StmtSend(
+		builder.ExprMonetary(
+			builder.ExprAsset("USD/2"),
+			builder.ExprNumberBigInt(big.NewInt(42)),
+		),
+		builder.SrcInorder(
+			builder.SrcAccount(
+				builder.ExprAccount("src1"),
+			),
+			builder.SrcAccount(
+				builder.ExprAccount("src2"),
+			),
+		),
+		builder.DestAccount(
+			builder.ExprAccount("dest"),
+		),
+	)
+
+	vars, _, script := builder.BuildProgram(stmt)
+	snaps.MatchInlineSnapshot(t, script, snaps.Inline(`vars {
+  account $account_0
+  account $account_1
+  account $account_2
+  asset $asset_0
+}
+
+send [$asset_0 42] (
+  source = {
+    $account_0
+    $account_1
+  }
+  destination = $account_2
+)`))
+
+	require.Equal(t, map[string]string{
+		"account_0": "src1",
+		"account_1": "src2",
+		"account_2": "dest",
+		"asset_0":   "USD/2",
+	}, vars)
+}
+
+func TestInorderNested(t *testing.T) {
+	stmt := builder.StmtSend(
+		builder.ExprMonetary(
+			builder.ExprAsset("USD/2"),
+			builder.ExprNumberBigInt(big.NewInt(42)),
+		),
+		builder.SrcInorder(
+			builder.SrcAccount(
+				builder.ExprAccount("src1"),
+			),
+			builder.SrcAccount(
+				builder.ExprAccount("src2"),
+			),
+			builder.SrcInorder(
+				builder.SrcAccount(
+					builder.ExprAccount("src_nested1"),
+				),
+				builder.SrcAccount(
+					builder.ExprAccount("src_nested2"),
+				),
+			),
+			builder.SrcAccount(
+				builder.ExprAccount("src_upper"),
+			),
+		),
+		builder.DestAccount(
+			builder.ExprAccount("dest"),
+		),
+	)
+
+	_, _, script := builder.BuildProgram(stmt)
+	snaps.MatchInlineSnapshot(t, script, snaps.Inline(`vars {
+  account $account_0
+  account $account_1
+  account $account_2
+  account $account_3
+  account $account_4
+  account $account_5
+  asset $asset_0
+}
+
+send [$asset_0 42] (
+  source = {
+    $account_0
+    $account_1
+    {
+      $account_2
+      $account_3
+    }
+    $account_4
+  }
+  destination = $account_5
+)`))
+}
+
+func TestInorderWithColors(t *testing.T) {
+	stmt := builder.StmtSend(
+		builder.ExprMonetary(
+			builder.ExprAsset("USD/2"),
+			builder.ExprNumberBigInt(big.NewInt(42)),
+		),
+		builder.SrcInorder(
+			builder.SrcColored(
+				builder.ExprAccount("acc"),
+				builder.ExprString("col"),
+			),
+			builder.SrcAccount(
+				builder.ExprAccount("src2"),
+			),
+		),
+		builder.DestAccount(
+			builder.ExprAccount("dest"),
+		),
+	)
+
+	_, _, script := builder.BuildProgram(stmt)
+	snaps.MatchInlineSnapshot(t, script, snaps.Inline(`vars {
+  account $account_0
+  account $account_1
+  account $account_2
+  string $string_0
+  asset $asset_0
+}
+
+send [$asset_0 42] (
+  source = {
+    $account_0 \ $string_0
+    $account_1
+  }
+  destination = $account_2
+)`))
+}
+
+func TestWithExternVar(t *testing.T) {
+	// The builder module exposes a type-safe API to create scripts
+
+	// We can create (typed) vars this way:
+	accVar := builder.NewAccountVar()
+	amtVar := builder.NewNumberVar()
+
+	// sources, destinations, expressions, and statements are typed, so you can never
+	// use a source node instead of a expression node, and so on.
+	// In addition to that, expressions are typed, so you can't use a string expression
+	// instead of a number expression
+	stmt := builder.StmtSend(
+		builder.ExprMonetary(
+			builder.ExprAsset("USD/2"),
+			builder.ExprVar(&amtVar), // <- you can reference vars this way (identified by address)
+		),
+		builder.SrcAccount(
+			builder.ExprVar(&accVar),
+		),
+		builder.DestAccount(
+			builder.ExprAccount("dest"),
+		),
+	)
+
+	// When you build the program, it'll create 3 values:
+	vars, varsEnv, script := builder.BuildProgram(stmt)
+
+	// 1: (vars) the map[string]string of KNOWN variables. This is generated via the
+	// strings literals you pass (in the example above, USD/2 and @dest).
+	// This way, you'll pass this map to the tx and numscript will handle interpolation instead of
+	// handling that in this lib
+	require.Equal(t, map[string]string{
+		"asset_0":   "USD/2",
+		"account_1": "dest",
+	}, vars)
+
+	// 2: (varsEnv) The env of the NAMES of each variable that is referenced within the script.
+	// You'll reference them by ptr address. The "Fill*()" methods are typed, and return you the name of the var,
+	// and the "stringified" value of the var content (in the case of account/asset/string, the string itself)
+	// Behaviour of Fill*() of vars that are never referenced in the script (thus, whose name is never allocated) is undefined
+	// (it may panic in the future)
+	//
+	// user code would likely be something like:
+	//
+	// varsCp := maps.Clone(vars)
+	// k, v := varsEnv.FillAccount(..)
+	// varsCp[k] = v
+	// (etc)
+	k, v := varsEnv.FillAccount(&accVar, "my_src")
+	require.Equal(t, "account_0", k)
+	require.Equal(t, "my_src", v)
+
+	k, v = varsEnv.FillNumber(&amtVar, big.NewInt(42))
+	require.Equal(t, "number_0", k)
+	require.Equal(t, "42", v)
+
+	// 3: (script) the script itself. It's generated while trying to be as "stable" as possible,
+	// e.g. vars are declared with an hardcoded order of types, and with an increasing order of names
+	snaps.MatchInlineSnapshot(t, script, snaps.Inline(`vars {
+  account $account_0
+  account $account_1
+  asset $asset_0
+  number $number_0
+}
+
+send [$asset_0 $number_0] (
+  source = $account_0
+  destination = $account_1
+)`))
+}