feat: color as first class citizen

inputs.schema.json

@@ -21,18 +21,31 @@
   },
   "definitions": {
     "Balances": {
+      "type": "array",
+      "description": "List of account balances. The (account, asset, color) triple of each entry must be unique within the list.",
+      "items": { "$ref": "#/definitions/BalanceRow" }
+    },
+
+    "BalanceRow": {
       "type": "object",
-      "description": "Map of account names to asset balances",
+      "description": "The balance of a given (account, asset, color)",
       "additionalProperties": false,
-      "patternProperties": {
-        "^([a-zA-Z0-9_-]+(:[a-zA-Z0-9_-]+)*)$": {
-          "type": "object",
-          "additionalProperties": false,
-          "patternProperties": {
-            "^([A-Z]+(/[0-9]+)?)$": {
-              "type": "number"
-            }
-          }
+      "required": ["account", "asset", "amount"],
+      "properties": {
+        "account": {
+          "type": "string",
+          "pattern": "^([a-zA-Z0-9_-]+(:[a-zA-Z0-9_-]+)*)$"
+        },
+        "asset": {
+          "type": "string",
+          "pattern": "^([A-Z]+(/[0-9]+)?)$"
+        },
+        "amount": {
+          "type": "number"
+        },
+        "color": {
+          "type": "string",
+          "pattern": "^[A-Z]*$"
         }
       }
     },

internal/analysis/check.go

@@ -690,7 +690,6 @@ func (res *CheckResult) checkSource(source parser.Source) {
 	case *parser.SourceWithScaling:
 		res.checkExpression(source.Address, TypeAccount)
 		res.checkExpression(source.Through, TypeAccount)
-		res.checkExpression(source.Color, TypeString)
 
 	case *parser.SourceInorder:
 		for _, source := range source.Sources {

internal/cmd/run.go

@@ -57,6 +57,16 @@ func run(scriptPath string, opts RunArgs) error {
 		return fmt.Errorf("failed to parse inputs file '%s' as JSON: %w", inputsPath, err)
 	}
 
+	// Reject a malformed inputs file before running anything: a balance list is a
+	// map keyed by (account, asset, color), so a repeated key is ambiguous.
+	if dup, ok := inputs.Balances.FirstDuplicate(); ok {
+		key := fmt.Sprintf("account=%q asset=%q", dup.Account, dup.Asset)
+		if dup.Color != "" {
+			key += fmt.Sprintf(" color=%q", dup.Color)
+		}
+		return fmt.Errorf("invalid inputs file '%s': balances must not contain duplicate entries: duplicate entry for %s", inputsPath, key)
+	}
+
 	featureFlags := map[string]struct{}{}
 	for _, flag := range inputs.FeatureFlags {
 		featureFlags[flag] = struct{}{}

internal/cmd/test_init.go

@@ -11,7 +11,6 @@ import (
 	"github.com/formancehq/numscript/internal/interpreter"
 	"github.com/formancehq/numscript/internal/parser"
 	"github.com/formancehq/numscript/internal/specs_format"
-	"github.com/formancehq/numscript/internal/utils"
 	"github.com/spf13/cobra"
 )
 
@@ -114,10 +113,12 @@ func makeSpecsFile(
 	defaultBalance *big.Int,
 ) (specs_format.Specs, error) {
 
-	store := TestInitStore{
+	store := &TestInitStore{
 		DefaultBalance: defaultBalance,
-		Balances:       make(interpreter.Balances),
-		Meta:           make(interpreter.AccountsMetadata),
+		StaticStore: interpreter.StaticStore{
+			Balances: interpreter.Balances{},
+			Meta:     make(interpreter.AccountsMetadata),
+		},
 	}
 
 	res, iErr := interpreter.RunProgram(
@@ -147,7 +148,7 @@ func makeSpecsFile(
 				program,
 				vars,
 				featureFlags,
-				&missingFundsErr.Needed,
+				defaultBalance,
 			)
 		}
 
@@ -161,7 +162,7 @@ func makeSpecsFile(
 
 	specs := specs_format.Specs{
 		Schema:       "https://raw.githubusercontent.com/formancehq/numscript/main/specs.schema.json",
-		Balances:     store.Balances,
+		Balances:     store.StaticStore.Balances,
 		Vars:         vars,
 		FeatureFlags: featureFlags_,
 		TestCases: []specs_format.TestCase{
@@ -200,40 +201,52 @@ func runTestInitCmd(opts testInitArgs) error {
 
 type TestInitStore struct {
 	DefaultBalance *big.Int
-	Balances       interpreter.Balances
-	Meta           interpreter.AccountsMetadata
+	StaticStore    interpreter.StaticStore
 }
 
-func (s TestInitStore) GetBalances(_ context.Context, q interpreter.BalanceQuery) (interpreter.Balances, error) {
-	outputBalance := interpreter.Balances{}
-	for queriedAccount, queriedCurrencies := range q {
-
-		for _, curr := range queriedCurrencies {
-			amt := utils.NestedMapGetOrPutDefault(s.Balances, queriedAccount, curr, func() *big.Int {
-				return new(big.Int).Set(s.DefaultBalance)
-			})
+func (s *TestInitStore) GetBalances(ctx context.Context, q interpreter.BalanceQuery) (interpreter.Balances, error) {
+	balances, err := s.StaticStore.GetBalances(ctx, q)
+	if err != nil {
+		return nil, err
+	}
 
-			outputAccountBalance := utils.MapGetOrPutDefault(outputBalance, queriedAccount, func() interpreter.AccountBalance {
-				return interpreter.AccountBalance{}
-			})
+	type key struct{ account, asset, color string }
 
-			outputAccountBalance[curr] = new(big.Int).Set(amt)
-		}
+	// StaticStore.GetBalances materializes a zero-amount row for every queried
+	// (account, asset, color), so its output can't tell a known account from an
+	// unknown one. Track what we've actually funded ourselves instead.
+	stored := make(map[key]struct{}, len(s.StaticStore.Balances))
+	for _, b := range s.StaticStore.Balances {
+		stored[key{b.Account, b.Asset, b.Color}] = struct{}{}
 	}
 
-	return outputBalance, nil
-}
+	for i := range balances {
+		b := &balances[i]
+		k := key{b.Account, b.Asset, b.Color}
+		if _, ok := stored[k]; ok {
+			continue
+		}
 
-func (s TestInitStore) GetAccountsMetadata(c context.Context, q interpreter.MetadataQuery) (interpreter.AccountsMetadata, error) {
-	outputMeta := interpreter.AccountsMetadata{}
-	for queriedAccount, queriedCurrencies := range q {
-		for _, curr := range queriedCurrencies {
-			outputAccountMeta := utils.MapGetOrPutDefault(outputMeta, queriedAccount, func() interpreter.AccountMetadata {
-				return interpreter.AccountMetadata{}
-			})
-			outputAccountMeta[curr] = ""
+		// Unknown (account, asset, color): fund it with the default balance, and
+		// remember it so later queries (and the generated specs file) see it.
+		amount := new(big.Int)
+		if s.DefaultBalance != nil {
+			amount.Set(s.DefaultBalance)
 		}
+		b.Amount = amount
+
+		s.StaticStore.Balances = append(s.StaticStore.Balances, interpreter.BalanceRow{
+			Account: b.Account,
+			Asset:   b.Asset,
+			Color:   b.Color,
+			Amount:  new(big.Int).Set(amount),
+		})
+		stored[k] = struct{}{}
 	}
 
-	return outputMeta, nil
+	return balances, nil
+}
+
+func (s *TestInitStore) GetAccountsMetadata(c context.Context, q interpreter.MetadataQuery) (interpreter.AccountsMetadata, error) {
+	return s.StaticStore.GetAccountsMetadata(c, q)
 }

internal/cmd/test_init_test.go

@@ -19,9 +19,7 @@ func TestMakeSpecsFileRetryForMissingFunds(t *testing.T) {
 
 	require.Nil(t, err)
 	require.Equal(t, interpreter.Balances{
-		"alice": interpreter.AccountBalance{
-			"USD/2": big.NewInt(10000),
-		},
+		{Account: "alice", Asset: "USD/2", Amount: big.NewInt(10000)},
 	}, out.Balances)
 }
 

internal/interpreter/accounts_metadata.go

@@ -4,6 +4,9 @@ import (
 	"github.com/formancehq/numscript/internal/utils"
 )
 
+type AccountMetadata = map[string]string
+type AccountsMetadata map[string]AccountMetadata
+
 func (m AccountsMetadata) fetchAccountMetadata(account string) AccountMetadata {
 	return utils.MapGetOrPutDefault(m, account, func() AccountMetadata {
 		return AccountMetadata{}

internal/interpreter/args_parser.go

@@ -14,19 +14,21 @@ func NewArgsParser(args []Value) *argsParser {
 	}
 }
 
-func parseArg[T any](p *argsParser, r parser.Range, expect func(Value, parser.Range) (*T, InterpreterError)) *T {
+func parseArg[T any](p *argsParser, r parser.Range, expect func(Value, parser.Range) (T, InterpreterError)) T {
 	index := p.parsedArgsCount
 	p.parsedArgsCount++
 
 	if p.err != nil || index >= len(p.args) {
-		return nil
+		var default_ T
+		return default_
 	}
 
 	arg := p.args[index]
 	parsed, err := expect(arg, r)
 	if err != nil {
 		p.err = err
-		return nil
+		var default_ T
+		return default_
 	}
 	return parsed
 }

internal/interpreter/args_parser_test.go

@@ -47,8 +47,8 @@ func TestParseValid(t *testing.T) {
 	require.NotNil(t, a1, "a1 should not be nil")
 	require.NotNil(t, a2, "a2 should not be nil")
 
-	require.Equal(t, *a1, *big.NewInt(42))
-	require.Equal(t, *a2, "user:001")
+	require.Equal(t, a1, MonetaryInt(*big.NewInt(42)))
+	require.Equal(t, a2, AccountAddress("user:001"))
 }
 
 func TestParseBadType(t *testing.T) {

internal/interpreter/asset_scaling.go

@@ -4,18 +4,18 @@ import (
 	"fmt"
 	"math/big"
 	"slices"
-	"strconv"
 	"strings"
 
 	"github.com/formancehq/numscript/internal/utils"
 )
 
-func assetToScaledAsset(asset string) string {
-	parts := strings.Split(asset, "/")
+func assetToScaledAsset(asset Asset) Asset {
+	strAsset := string(asset)
+	parts := strings.Split(strAsset, "/")
 	if len(parts) == 1 {
-		return asset + "/*"
+		return Asset(strAsset + "/*")
 	}
-	return parts[0] + "/*"
+	return Asset(parts[0] + "/*")
 }
 
 func buildScaledAsset(baseAsset string, scale int64) string {
@@ -25,25 +25,17 @@ func buildScaledAsset(baseAsset string, scale int64) string {
 	return fmt.Sprintf("%s/%d", baseAsset, scale)
 }
 
-func getAssetScale(asset string) (string, int64) {
-	parts := strings.Split(asset, "/")
-	if len(parts) == 2 {
-		scale, err := strconv.ParseInt(parts[1], 10, 64)
-		if err == nil {
-			return parts[0], scale
-		}
-		// fallback if parsing fails
-		return parts[0], 0
-	}
-	return asset, 0
-}
-
-func getAssets(balance AccountBalance, baseAsset string) map[int64]*big.Int {
+func getAssets(accountBalances []AccountBalance, baseAsset string) map[int64]*big.Int {
 	result := make(map[int64]*big.Int)
-	for asset, amount := range balance {
-		if strings.HasPrefix(asset, baseAsset) {
-			_, scale := getAssetScale(asset)
-			result[scale] = amount
+	for _, accBalance := range accountBalances {
+		if accBalance.Color != "" {
+			// scaling converts only uncolored balances, and emits uncolored
+			// postings, so a colored balance must not be treated as available
+			continue
+		}
+		accBalanceAsset, scale := Asset(accBalance.Asset).GetBaseAndScale()
+		if accBalanceAsset == baseAsset {
+			result[scale] = new(big.Int).Set(accBalance.Amount)
 		}
 	}
 	return result

internal/interpreter/asset_scaling_test.go

@@ -7,6 +7,22 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestGetAssetsExcludesColoredBalances(t *testing.T) {
+	// Scaling converts only uncolored balances and emits uncolored postings, so a
+	// colored balance for the same base asset must not be offered as a candidate.
+	assets := getAssets([]AccountBalance{
+		{Asset: "USD", Color: "", Amount: big.NewInt(2)},
+		{Asset: "USD/4", Color: "RED", Amount: big.NewInt(999)}, // colored: excluded
+		{Asset: "USD/2", Color: "", Amount: big.NewInt(50)},
+		{Asset: "EUR", Color: "", Amount: big.NewInt(7)}, // different base: excluded
+	}, "USD")
+
+	require.Equal(t, map[int64]*big.Int{
+		0: big.NewInt(2),
+		2: big.NewInt(50),
+	}, assets)
+}
+
 func TestScalingAvoidSwappingAlreadyHaveAsset(t *testing.T) {
 	// Need [USD/2 200]
 	// Got: {USD/2 100, USD 2}