diff --git a/inputs.schema.json b/inputs.schema.json index ee73b798..d607d4e1 100644 --- a/inputs.schema.json +++ b/inputs.schema.json @@ -30,13 +30,31 @@ "additionalProperties": false, "patternProperties": { "^([A-Z]+(/[0-9]+)?)$": { - "type": "number" + "oneOf": [ + { "type": "number", "description": "Uncolored shorthand" }, + { "$ref": "#/definitions/BalanceEntry" }, + { "type": "array", "items": { "$ref": "#/definitions/BalanceEntry" } } + ] } } } } }, + "BalanceEntry": { + "type": "object", + "additionalProperties": false, + "required": ["amount"], + "properties": { + "color": { + "type": "string", + "pattern": "^[A-Z]*$", + "description": "Color bucket. Omit (or set to \"\") for the uncolored bucket." + }, + "amount": { "type": "number" } + } + }, + "VariablesMap": { "type": "object", "description": "Map of variable name to variable stringified value", diff --git a/internal/cmd/test_init.go b/internal/cmd/test_init.go index 9027b2c6..d98d025f 100644 --- a/internal/cmd/test_init.go +++ b/internal/cmd/test_init.go @@ -206,18 +206,25 @@ type TestInitStore struct { 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 { + for queriedAccount, queriedItems := range q { + for _, item := range queriedItems { + accBalance := utils.MapGetOrPutDefault(s.Balances, queriedAccount, func() interpreter.AccountBalance { + return interpreter.AccountBalance{} + }) + colorMap := utils.MapGetOrPutDefault(accBalance, item.Asset, func() interpreter.ColorBalance { + return interpreter.ColorBalance{} + }) + amt := utils.MapGetOrPutDefault(colorMap, item.Color, func() *big.Int { return new(big.Int).Set(s.DefaultBalance) }) outputAccountBalance := utils.MapGetOrPutDefault(outputBalance, queriedAccount, func() interpreter.AccountBalance { return interpreter.AccountBalance{} }) - - outputAccountBalance[curr] = new(big.Int).Set(amt) + outputColorMap := utils.MapGetOrPutDefault(outputAccountBalance, item.Asset, func() interpreter.ColorBalance { + return interpreter.ColorBalance{} + }) + outputColorMap[item.Color] = new(big.Int).Set(amt) } } diff --git a/internal/cmd/test_init_test.go b/internal/cmd/test_init_test.go index f6fa865f..13b88aca 100644 --- a/internal/cmd/test_init_test.go +++ b/internal/cmd/test_init_test.go @@ -20,7 +20,7 @@ func TestMakeSpecsFileRetryForMissingFunds(t *testing.T) { require.Nil(t, err) require.Equal(t, interpreter.Balances{ "alice": interpreter.AccountBalance{ - "USD/2": big.NewInt(10000), + "USD/2": interpreter.Uncolored(big.NewInt(10000)), }, }, out.Balances) } diff --git a/internal/interpreter/__snapshots__/balances_test.snap b/internal/interpreter/__snapshots__/balances_test.snap index b3ac285e..7c7b89b1 100755 --- a/internal/interpreter/__snapshots__/balances_test.snap +++ b/internal/interpreter/__snapshots__/balances_test.snap @@ -1,7 +1,7 @@ [TestPrettyPrintBalance - 1] -| Account | Asset  | Balance | -| alice | EUR/2 | 1 | -| alice | USD/1234 | 999999 | -| bob | BTC | 3 | +| Account | Asset  | Color | Balance | +| alice | EUR/2 | | 1 | +| alice | USD/1234 | | 999999 | +| bob | BTC | | 3 | --- diff --git a/internal/interpreter/asset_scaling.go b/internal/interpreter/asset_scaling.go index 48c5fd41..12a46569 100644 --- a/internal/interpreter/asset_scaling.go +++ b/internal/interpreter/asset_scaling.go @@ -38,13 +38,21 @@ func getAssetScale(asset string) (string, int64) { return asset, 0 } +// getAssets returns, for a given baseAsset, the per-scale uncolored balance. +// Asset scaling operates on the uncolored bucket only — colored funds are not +// implicitly converted across scales. func getAssets(balance 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 asset, colorMap := range balance { + if !strings.HasPrefix(asset, baseAsset) { + continue + } + amount, ok := colorMap[""] + if !ok { + continue } + _, scale := getAssetScale(asset) + result[scale] = amount } return result } diff --git a/internal/interpreter/balances.go b/internal/interpreter/balances.go index 36f63c61..111535e3 100644 --- a/internal/interpreter/balances.go +++ b/internal/interpreter/balances.go @@ -2,110 +2,124 @@ package interpreter import ( "math/big" - "strings" "github.com/formancehq/numscript/internal/utils" ) +// Uncolored builds a ColorBalance holding a single amount under the empty +// color key (the "no color" bucket). It is meant to keep test setup terse. +func Uncolored(amount *big.Int) ColorBalance { + return ColorBalance{"": amount} +} + func (b Balances) DeepClone() Balances { cloned := make(Balances) for account, accountBalances := range b { - for asset, amount := range accountBalances { - utils.NestedMapGetOrPutDefault(cloned, account, asset, func() *big.Int { - return new(big.Int).Set(amount) - }) + clonedAcc := AccountBalance{} + cloned[account] = clonedAcc + for asset, colorMap := range accountBalances { + clonedColors := ColorBalance{} + clonedAcc[asset] = clonedColors + for color, amount := range colorMap { + clonedColors[color] = new(big.Int).Set(amount) + } } } return cloned } -func coloredAsset(asset string, color *string) string { - if color == nil || *color == "" { - return asset - } - - // note: 1 <= len(parts) <= 2 - parts := strings.Split(asset, "/") - - coloredAsset := parts[0] + "_" + *color - if len(parts) > 1 { - coloredAsset += "/" + parts[1] - } - return coloredAsset -} - -// Get the (account, asset) tuple from the Balances -// if the tuple is not present, it will write a big.NewInt(0) in it and return it -func (b Balances) fetchBalance(account string, uncoloredAsset string, color string) *big.Int { - return utils.NestedMapGetOrPutDefault(b, account, coloredAsset(uncoloredAsset, &color), func() *big.Int { +// Get the (account, asset, color) balance from Balances. +// If the entry is not present, it will write a big.NewInt(0) in it and return it. +func (b Balances) fetchBalance(account string, asset string, color string) *big.Int { + accBalance := utils.MapGetOrPutDefault(b, account, func() AccountBalance { + return AccountBalance{} + }) + colorMap := utils.MapGetOrPutDefault(accBalance, asset, func() ColorBalance { + return ColorBalance{} + }) + return utils.MapGetOrPutDefault(colorMap, color, func() *big.Int { return new(big.Int) }) } -func (b Balances) has(account string, asset string) bool { - accountBalances := utils.MapGetOrPutDefault(b, account, func() AccountBalance { - return AccountBalance{} - }) - - _, ok := accountBalances[asset] +func (b Balances) has(account string, asset string, color string) bool { + accountBalances, ok := b[account] + if !ok { + return false + } + colorMap, ok := accountBalances[asset] + if !ok { + return false + } + _, ok = colorMap[color] return ok } -// given a BalanceQuery, return a new query which only contains needed (asset, account) pairs +// given a BalanceQuery, return a new query which only contains needed (asset, color) pairs // (that is, the ones that aren't already cached) func (b Balances) filterQuery(q BalanceQuery) BalanceQuery { filteredQuery := BalanceQuery{} - for accountName, queriedCurrencies := range q { - filteredCurrencies := utils.Filter(queriedCurrencies, func(currency string) bool { - return !b.has(accountName, currency) + for accountName, queriedItems := range q { + filteredItems := utils.Filter(queriedItems, func(item AssetColor) bool { + return !b.has(accountName, item.Asset, item.Color) }) - if len(filteredCurrencies) > 0 { - filteredQuery[accountName] = filteredCurrencies + if len(filteredItems) > 0 { + filteredQuery[accountName] = filteredItems } - } return filteredQuery } // Merge balances by adding balances in the "update" arg func (b Balances) Merge(update Balances) { - // merge queried balance for acc, accBalances := range update { cachedAcc := utils.MapGetOrPutDefault(b, acc, func() AccountBalance { return AccountBalance{} }) - for curr, amt := range accBalances { - cachedAcc[curr] = amt + for asset, colorMap := range accBalances { + cachedColors := utils.MapGetOrPutDefault(cachedAcc, asset, func() ColorBalance { + return ColorBalance{} + }) + for color, amt := range colorMap { + cachedColors[color] = amt + } } } } func (b Balances) PrettyPrint() string { - header := []string{"Account", "Asset", "Balance"} + header := []string{"Account", "Asset", "Color", "Balance"} var rows [][]string for account, accBalances := range b { - for asset, balance := range accBalances { - row := []string{account, asset, balance.String()} - rows = append(rows, row) + for asset, colorMap := range accBalances { + for color, balance := range colorMap { + rows = append(rows, []string{account, asset, color, balance.String()}) + } } } return utils.CsvPretty(header, rows, true) } func CompareBalances(b1 Balances, b2 Balances) bool { - return utils.Map2Cmp(b1, b2, func(ab1, ab2 *big.Int) bool { - return ab1.Cmp(ab2) == 0 + return utils.MapCmp(b1, b2, func(ab1, ab2 AccountBalance) bool { + return utils.MapCmp(ab1, ab2, func(cm1, cm2 ColorBalance) bool { + return utils.MapCmp(cm1, cm2, func(a1, a2 *big.Int) bool { + return a1.Cmp(a2) == 0 + }) + }) }) } // Returns whether the first value is a subset of the second one func CompareBalancesIncluding(b1 Balances, b2 Balances) bool { return utils.MapIncludes(b2, b1, func(a2 AccountBalance, a1 AccountBalance) bool { - return utils.MapIncludes(a2, a1, func(a2 *big.Int, a1 *big.Int) bool { - return a2.Cmp(a1) == 0 + return utils.MapIncludes(a2, a1, func(cm2, cm1 ColorBalance) bool { + return utils.MapIncludes(cm2, cm1, func(x2, x1 *big.Int) bool { + return x2.Cmp(x1) == 0 + }) }) }) } diff --git a/internal/interpreter/balances_json.go b/internal/interpreter/balances_json.go new file mode 100644 index 00000000..77551265 --- /dev/null +++ b/internal/interpreter/balances_json.go @@ -0,0 +1,142 @@ +package interpreter + +import ( + "bytes" + "encoding/json" + "fmt" + "math/big" + "sort" +) + +// Balances JSON shape. +// +// Canonical write form per (account, asset) entry: +// +// - a bare integer when only the uncolored bucket is set: +// "USD/2": 100 +// +// - an array of value-objects when one or more colors are present: +// "USD/2": [ +// { "amount": 100 }, // uncolored — no "color" field +// { "color": "RED", "amount": 50 } // colored +// ] +// +// Tolerant read accepts three forms per (account, asset) entry: +// 1. JSON number → single uncolored bucket +// 2. JSON object → single value-object with optional "color" and required "amount" +// 3. JSON array → list of value-objects (canonical multi-color form) +// +// Adding orthogonal dimensions in the future (scopes, scales, …) is done by +// extending the value-object schema rather than adding a new nesting level. + +type balanceEntry struct { + Color string `json:"color,omitempty"` + Amount *big.Int `json:"amount"` +} + +func (b Balances) MarshalJSON() ([]byte, error) { + type assetWire = json.RawMessage + + out := make(map[string]map[string]assetWire, len(b)) + for account, accBalances := range b { + assets := make(map[string]assetWire, len(accBalances)) + for asset, colorMap := range accBalances { + raw, err := marshalColorBalance(colorMap) + if err != nil { + return nil, fmt.Errorf("balances[%q][%q]: %w", account, asset, err) + } + assets[asset] = raw + } + out[account] = assets + } + return json.Marshal(out) +} + +func marshalColorBalance(cb ColorBalance) (json.RawMessage, error) { + if len(cb) == 1 { + if amount, ok := cb[""]; ok { + return json.Marshal(amount) + } + } + + colors := make([]string, 0, len(cb)) + for c := range cb { + colors = append(colors, c) + } + sort.Strings(colors) + + entries := make([]balanceEntry, len(colors)) + for i, c := range colors { + entries[i] = balanceEntry{Color: c, Amount: cb[c]} + } + return json.Marshal(entries) +} + +func (b *Balances) UnmarshalJSON(data []byte) error { + var raw map[string]map[string]json.RawMessage + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + + out := make(Balances, len(raw)) + for account, assets := range raw { + accBalance := make(AccountBalance, len(assets)) + for asset, rawValue := range assets { + cb, err := unmarshalColorBalance(rawValue) + if err != nil { + return fmt.Errorf("balances[%q][%q]: %w", account, asset, err) + } + accBalance[asset] = cb + } + out[account] = accBalance + } + *b = out + return nil +} + +func unmarshalColorBalance(data []byte) (ColorBalance, error) { + trimmed := bytes.TrimSpace(data) + if len(trimmed) == 0 { + return nil, fmt.Errorf("empty balance entry") + } + + switch trimmed[0] { + case '[': + var entries []balanceEntry + dec := json.NewDecoder(bytes.NewReader(trimmed)) + dec.DisallowUnknownFields() + if err := dec.Decode(&entries); err != nil { + return nil, err + } + cb := make(ColorBalance, len(entries)) + for i, e := range entries { + if e.Amount == nil { + return nil, fmt.Errorf("entry %d: missing \"amount\"", i) + } + if _, exists := cb[e.Color]; exists { + return nil, fmt.Errorf("entry %d: duplicate color %q", i, e.Color) + } + cb[e.Color] = e.Amount + } + return cb, nil + + case '{': + var entry balanceEntry + dec := json.NewDecoder(bytes.NewReader(trimmed)) + dec.DisallowUnknownFields() + if err := dec.Decode(&entry); err != nil { + return nil, err + } + if entry.Amount == nil { + return nil, fmt.Errorf("missing \"amount\" field") + } + return ColorBalance{entry.Color: entry.Amount}, nil + + default: + amount := new(big.Int) + if err := json.Unmarshal(trimmed, amount); err != nil { + return nil, fmt.Errorf("expected number, value-object, or array; got %s", string(trimmed)) + } + return ColorBalance{"": amount}, nil + } +} diff --git a/internal/interpreter/balances_test.go b/internal/interpreter/balances_test.go index d3d094be..6e73312f 100644 --- a/internal/interpreter/balances_test.go +++ b/internal/interpreter/balances_test.go @@ -11,52 +11,96 @@ import ( func TestFilterQuery(t *testing.T) { fullBalance := Balances{ "alice": AccountBalance{ - "EUR/2": big.NewInt(1), - "USD/2": big.NewInt(2), + "EUR/2": Uncolored(big.NewInt(1)), + "USD/2": Uncolored(big.NewInt(2)), }, "bob": AccountBalance{ - "BTC": big.NewInt(3), + "BTC": Uncolored(big.NewInt(3)), }, } filteredQuery := fullBalance.filterQuery(BalanceQuery{ - "alice": []string{"GBP/2", "YEN", "EUR/2"}, - "bob": []string{"BTC"}, - "charlie": []string{"ETH"}, + "alice": []AssetColor{{Asset: "GBP/2"}, {Asset: "YEN"}, {Asset: "EUR/2"}}, + "bob": []AssetColor{{Asset: "BTC"}}, + "charlie": []AssetColor{{Asset: "ETH"}}, }) require.Equal(t, BalanceQuery{ - "alice": []string{"GBP/2", "YEN"}, - "charlie": []string{"ETH"}, + "alice": []AssetColor{{Asset: "GBP/2"}, {Asset: "YEN"}}, + "charlie": []AssetColor{{Asset: "ETH"}}, + }, filteredQuery) +} + +func TestFilterQueryDistinguishesColors(t *testing.T) { + t.Parallel() + + fullBalance := Balances{ + "alice": AccountBalance{ + "USD/2": ColorBalance{ + "": big.NewInt(1), + "GRANTS": big.NewInt(2), + }, + }, + } + + filteredQuery := fullBalance.filterQuery(BalanceQuery{ + "alice": []AssetColor{ + {Asset: "USD/2"}, + {Asset: "USD/2", Color: "GRANTS"}, + {Asset: "USD/2", Color: "OPS"}, + }, + }) + + require.Equal(t, BalanceQuery{ + "alice": []AssetColor{{Asset: "USD/2", Color: "OPS"}}, }, filteredQuery) } func TestCloneBalances(t *testing.T) { fullBalance := Balances{ "alice": AccountBalance{ - "EUR/2": big.NewInt(1), - "USD/2": big.NewInt(2), + "EUR/2": Uncolored(big.NewInt(1)), + "USD/2": Uncolored(big.NewInt(2)), }, "bob": AccountBalance{ - "BTC": big.NewInt(3), + "BTC": Uncolored(big.NewInt(3)), }, } cloned := fullBalance.DeepClone() - fullBalance["alice"]["USD/2"].Set(big.NewInt(42)) + fullBalance["alice"]["USD/2"][""].Set(big.NewInt(42)) - require.Equal(t, big.NewInt(2), cloned["alice"]["USD/2"]) + require.Equal(t, big.NewInt(2), cloned["alice"]["USD/2"][""]) +} + +func TestCloneBalancesPreservesColors(t *testing.T) { + t.Parallel() + + fullBalance := Balances{ + "alice": AccountBalance{ + "USD/2": ColorBalance{ + "RED": big.NewInt(10), + "BLUE": big.NewInt(20), + }, + }, + } + + cloned := fullBalance.DeepClone() + fullBalance["alice"]["USD/2"]["RED"].Set(big.NewInt(999)) + + require.Equal(t, big.NewInt(10), cloned["alice"]["USD/2"]["RED"]) + require.Equal(t, big.NewInt(20), cloned["alice"]["USD/2"]["BLUE"]) } func TestPrettyPrintBalance(t *testing.T) { fullBalance := Balances{ "alice": AccountBalance{ - "EUR/2": big.NewInt(1), - "USD/1234": big.NewInt(999999), + "EUR/2": Uncolored(big.NewInt(1)), + "USD/1234": Uncolored(big.NewInt(999999)), }, "bob": AccountBalance{ - "BTC": big.NewInt(3), + "BTC": Uncolored(big.NewInt(3)), }, } @@ -67,34 +111,53 @@ func TestCmpMaps(t *testing.T) { b1 := Balances{ "alice": AccountBalance{ - "EUR": big.NewInt(100), + "EUR": Uncolored(big.NewInt(100)), }, } b2 := Balances{ "alice": AccountBalance{ - "EUR": big.NewInt(42), + "EUR": Uncolored(big.NewInt(42)), }, } require.Equal(t, false, CompareBalances(b1, b2)) } +func TestCmpMapsDistinguishesColors(t *testing.T) { + t.Parallel() + + b1 := Balances{ + "alice": AccountBalance{ + "USD/2": ColorBalance{"RED": big.NewInt(100)}, + }, + } + + b2 := Balances{ + "alice": AccountBalance{ + "USD/2": ColorBalance{"BLUE": big.NewInt(100)}, + }, + } + + require.False(t, CompareBalances(b1, b2), + "same asset but different colors must not compare equal") +} + func TestCmpMapsIncluding(t *testing.T) { t.Run("including (subset)", func(t *testing.T) { b2 := Balances{ "alice": AccountBalance{ - "EUR": big.NewInt(100), + "EUR": Uncolored(big.NewInt(100)), }, "bob": AccountBalance{ - "EUR": big.NewInt(100), + "EUR": Uncolored(big.NewInt(100)), }, } b1 := Balances{ "alice": AccountBalance{ - "EUR": big.NewInt(100), + "EUR": Uncolored(big.NewInt(100)), }, } @@ -104,16 +167,16 @@ func TestCmpMapsIncluding(t *testing.T) { t.Run("different value", func(t *testing.T) { b2 := Balances{ "alice": AccountBalance{ - "EUR": big.NewInt(100), + "EUR": Uncolored(big.NewInt(100)), }, "bob": AccountBalance{ - "EUR": big.NewInt(100), + "EUR": Uncolored(big.NewInt(100)), }, } b1 := Balances{ "alice": AccountBalance{ - "EUR": big.NewInt(0), + "EUR": Uncolored(big.NewInt(0)), }, } @@ -123,23 +186,112 @@ func TestCmpMapsIncluding(t *testing.T) { t.Run("extra value", func(t *testing.T) { b2 := Balances{ "alice": AccountBalance{ - "EUR": big.NewInt(100), + "EUR": Uncolored(big.NewInt(100)), }, "bob": AccountBalance{ - "EUR": big.NewInt(100), + "EUR": Uncolored(big.NewInt(100)), }, } b1 := Balances{ "alice": AccountBalance{ - "EUR": big.NewInt(100), + "EUR": Uncolored(big.NewInt(100)), }, "extra-value": AccountBalance{ - "EUR": big.NewInt(100), + "EUR": Uncolored(big.NewInt(100)), }, } require.Equal(t, false, CompareBalancesIncluding(b1, b2)) }) + + t.Run("color-aware subset", func(t *testing.T) { + b2 := Balances{ + "alice": AccountBalance{ + "USD/2": ColorBalance{ + "": big.NewInt(100), + "RED": big.NewInt(50), + "BLUE": big.NewInt(25), + }, + }, + } + + b1 := Balances{ + "alice": AccountBalance{ + "USD/2": ColorBalance{"RED": big.NewInt(50)}, + }, + } + + require.True(t, CompareBalancesIncluding(b1, b2), + "colored subset of an asset must be considered included") + }) + + t.Run("missing color", func(t *testing.T) { + b2 := Balances{ + "alice": AccountBalance{ + "USD/2": ColorBalance{"": big.NewInt(100)}, + }, + } + + b1 := Balances{ + "alice": AccountBalance{ + "USD/2": ColorBalance{"RED": big.NewInt(50)}, + }, + } + + require.False(t, CompareBalancesIncluding(b1, b2), + "color present in subset but missing in superset must not be considered included") + }) +} + +func TestFetchBalanceCreatesEntriesLazily(t *testing.T) { + t.Parallel() + + b := Balances{} + got := b.fetchBalance("alice", "USD/2", "RED") + require.NotNil(t, got) + require.Equal(t, 0, got.Sign(), "freshly created entry should be zero") + + // mutating the returned big.Int should be reflected in subsequent fetches + got.SetInt64(42) + require.Equal(t, big.NewInt(42), b.fetchBalance("alice", "USD/2", "RED")) + + // distinct colors must be independent + require.Equal(t, big.NewInt(0), b.fetchBalance("alice", "USD/2", "BLUE")) +} + +func TestMergeBalancesIsColorAware(t *testing.T) { + t.Parallel() + + b := Balances{ + "alice": AccountBalance{ + "USD/2": ColorBalance{"RED": big.NewInt(1)}, + }, + } + + b.Merge(Balances{ + "alice": AccountBalance{ + "USD/2": ColorBalance{ + "RED": big.NewInt(99), // overwrites + "BLUE": big.NewInt(50), // adds new color under existing asset + }, + "EUR/2": Uncolored(big.NewInt(7)), // adds new asset + }, + "bob": AccountBalance{ + "BTC": Uncolored(big.NewInt(3)), // adds new account + }, + }) + + require.Equal(t, big.NewInt(99), b["alice"]["USD/2"]["RED"]) + require.Equal(t, big.NewInt(50), b["alice"]["USD/2"]["BLUE"]) + require.Equal(t, big.NewInt(7), b["alice"]["EUR/2"][""]) + require.Equal(t, big.NewInt(3), b["bob"]["BTC"][""]) +} + +func TestUncoloredHelper(t *testing.T) { + t.Parallel() + + got := Uncolored(big.NewInt(42)) + require.Equal(t, ColorBalance{"": big.NewInt(42)}, got) } diff --git a/internal/interpreter/batch_balances_query.go b/internal/interpreter/batch_balances_query.go index 141d17ab..dd3de3a6 100644 --- a/internal/interpreter/batch_balances_query.go +++ b/internal/interpreter/batch_balances_query.go @@ -54,11 +54,15 @@ func (st *programState) batchQuery(account string, asset string, color *string) if account == "world" { return } - asset = coloredAsset(asset, color) + colorStr := "" + if color != nil { + colorStr = *color + } + item := AssetColor{Asset: asset, Color: colorStr} previousValues := st.CurrentBalanceQuery[account] - if !slices.Contains(previousValues, asset) { - st.CurrentBalanceQuery[account] = append(previousValues, asset) + if !slices.Contains(previousValues, item) { + st.CurrentBalanceQuery[account] = append(previousValues, item) } } diff --git a/internal/interpreter/color_semantics_test.go b/internal/interpreter/color_semantics_test.go new file mode 100644 index 00000000..2c232f1a --- /dev/null +++ b/internal/interpreter/color_semantics_test.go @@ -0,0 +1,582 @@ +package interpreter_test + +import ( + "context" + "encoding/json" + "math/big" + "testing" + + "github.com/formancehq/numscript/internal/flags" + machine "github.com/formancehq/numscript/internal/interpreter" + "github.com/formancehq/numscript/internal/parser" + "github.com/stretchr/testify/require" +) + +// runColored is a small helper that parses a script, evaluates it with the +// experimental-asset-colors feature flag enabled, and returns the postings. +// All test cases below exercise the public color-of-money interface — the +// idea is to exercise it from many angles so any future regression in the +// numscript ↔ ledger contract is caught here, not downstream. +func runColored(t *testing.T, src string, store machine.StaticStore) ([]machine.Posting, error) { + t.Helper() + parsed := parser.Parse(src) + require.Empty(t, parsed.Errors, "unexpected parser errors: %v", parsed.Errors) + + result, err := machine.RunProgram( + context.Background(), + parsed.Value, + machine.VariablesMap{}, + store, + map[string]struct{}{flags.ExperimentalAssetColors: {}}, + ) + if err != nil { + return nil, err + } + return result.Postings, nil +} + +// Single colored send: a posting emitted from a "RED"-restricted source must +// carry Color="RED". +func TestColorSendPropagatesColor(t *testing.T) { + t.Parallel() + + src := ` + send [COIN 100] ( + source = @world \ "RED" + destination = @alice + ) + ` + store := machine.StaticStore{Balances: machine.Balances{}} + postings, err := runColored(t, src, store) + require.NoError(t, err) + + require.Equal(t, []machine.Posting{ + {Source: "world", Destination: "alice", Amount: big.NewInt(100), Asset: "COIN", Color: "RED"}, + }, postings) +} + +// Source-side color constraint: when a source restricts to a color, the +// emitted posting must carry that color and the funds must come from that +// exact bucket. +func TestColorSourceRestrictionEmitsColoredPosting(t *testing.T) { + t.Parallel() + + src := ` + send [COIN 20] ( + source = @acc \ "RED" + destination = @dest + ) + ` + store := machine.StaticStore{Balances: machine.Balances{ + "acc": machine.AccountBalance{ + "COIN": machine.ColorBalance{ + "": big.NewInt(1000), + "RED": big.NewInt(50), + }, + }, + }} + + postings, err := runColored(t, src, store) + require.NoError(t, err) + + require.Equal(t, []machine.Posting{ + {Source: "acc", Destination: "dest", Amount: big.NewInt(20), Asset: "COIN", Color: "RED"}, + }, postings) +} + +// Colored funds are strictly segregated: insufficient funds in a color must +// fail even when other colors (and the uncolored bucket) have plenty. +func TestColorIsolationRejectsSpendFromWrongColor(t *testing.T) { + t.Parallel() + + src := ` + send [COIN 100] ( + source = @acc \ "RED" + destination = @dest + ) + ` + store := machine.StaticStore{Balances: machine.Balances{ + "acc": machine.AccountBalance{ + "COIN": machine.ColorBalance{ + "": big.NewInt(10_000), + "BLUE": big.NewInt(10_000), + "RED": big.NewInt(20), + }, + }, + }} + + _, err := runColored(t, src, store) + require.Error(t, err) + var missing machine.MissingFundsErr + require.ErrorAs(t, err, &missing) + require.Equal(t, "COIN", missing.Asset) +} + +// "Color of money" is immutable: a posting emitted under color X stays +// color X end to end. We verify this by chaining sources and watching the +// emitted postings retain their original colors. +func TestColorImmutabilityThroughInorderSource(t *testing.T) { + t.Parallel() + + src := ` + send [COIN 150] ( + source = { + @acc \ "RED" + @acc \ "BLUE" + @acc + } + destination = @dest + ) + ` + store := machine.StaticStore{Balances: machine.Balances{ + "acc": machine.AccountBalance{ + "COIN": machine.ColorBalance{ + "": big.NewInt(100), + "BLUE": big.NewInt(30), + "RED": big.NewInt(20), + }, + }, + }} + + postings, err := runColored(t, src, store) + require.NoError(t, err) + + require.Equal(t, []machine.Posting{ + {Source: "acc", Destination: "dest", Amount: big.NewInt(20), Asset: "COIN", Color: "RED"}, + {Source: "acc", Destination: "dest", Amount: big.NewInt(30), Asset: "COIN", Color: "BLUE"}, + {Source: "acc", Destination: "dest", Amount: big.NewInt(100), Asset: "COIN", Color: ""}, + }, postings) +} + +// The uncolored bucket (Color="") is not pooled with any colored bucket — +// asking for 100 with color="" fails when only colored funds are available. +func TestUncoloredCannotDrawFromColoredFunds(t *testing.T) { + t.Parallel() + + src := ` + send [COIN 100] ( + source = @acc + destination = @dest + ) + ` + store := machine.StaticStore{Balances: machine.Balances{ + "acc": machine.AccountBalance{ + "COIN": machine.ColorBalance{ + "RED": big.NewInt(10_000), + }, + }, + }} + + _, err := runColored(t, src, store) + require.Error(t, err) + var missing machine.MissingFundsErr + require.ErrorAs(t, err, &missing) +} + +// Two adjacent colored postings with the same color from the same source +// should compact into a single posting (the funds queue logic). +func TestColoredPostingsCompactByColor(t *testing.T) { + t.Parallel() + + src := ` + send [COIN 30] ( + source = { + @acc \ "RED" + @acc \ "RED" + } + destination = @dest + ) + ` + store := machine.StaticStore{Balances: machine.Balances{ + "acc": machine.AccountBalance{ + "COIN": machine.ColorBalance{ + "RED": big.NewInt(100), + }, + }, + }} + + postings, err := runColored(t, src, store) + require.NoError(t, err) + require.Len(t, postings, 1) + require.Equal(t, "RED", postings[0].Color) + require.Equal(t, big.NewInt(30), postings[0].Amount) +} + +// Colored balance queries must include color in BalanceQuery items. +// We observe the store to verify the contract that gets sent across the +// numscript ↔ ledger boundary. +func TestBalanceQueryIncludesColor(t *testing.T) { + t.Parallel() + + store := machine.StaticStore{ + Balances: machine.Balances{ + "acc": machine.AccountBalance{ + "COIN": machine.ColorBalance{"RED": big.NewInt(1000)}, + }, + }, + } + + src := ` + send [COIN 100] ( + source = @acc \ "RED" + destination = @dest + ) + ` + parsed := parser.Parse(src) + require.Empty(t, parsed.Errors) + + // Wrap the store with an inline implementation that records GetBalances + // calls so we can assert what crosses the boundary. + var got []machine.BalanceQuery + spy := storeFunc{ + getBalances: func(ctx context.Context, q machine.BalanceQuery) (machine.Balances, error) { + cloned := machine.BalanceQuery{} + for acc, items := range q { + cloned[acc] = append([]machine.AssetColor(nil), items...) + } + got = append(got, cloned) + return store.GetBalances(ctx, q) + }, + getMetadata: store.GetAccountsMetadata, + } + + _, err := machine.RunProgram( + context.Background(), + parsed.Value, + machine.VariablesMap{}, + spy, + map[string]struct{}{flags.ExperimentalAssetColors: {}}, + ) + require.NoError(t, err) + + require.Len(t, got, 1, "expected exactly one batched balance query") + require.Equal(t, machine.BalanceQuery{ + "acc": {{Asset: "COIN", Color: "RED"}}, + }, got[0]) +} + +// storeFunc is a minimal machine.Store adapter built around function values. +type storeFunc struct { + getBalances func(context.Context, machine.BalanceQuery) (machine.Balances, error) + getMetadata func(context.Context, machine.MetadataQuery) (machine.AccountsMetadata, error) +} + +func (s storeFunc) GetBalances(ctx context.Context, q machine.BalanceQuery) (machine.Balances, error) { + return s.getBalances(ctx, q) +} + +func (s storeFunc) GetAccountsMetadata(ctx context.Context, q machine.MetadataQuery) (machine.AccountsMetadata, error) { + return s.getMetadata(ctx, q) +} + +// The legal color charset (^[A-Z]*$) is enforced — anything else must be +// rejected as a bad color literal. +func TestColorLiteralCharsetIsEnforced(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + color string + wantErr bool + }{ + {name: "uppercase ok", color: "RED", wantErr: false}, + {name: "empty ok (no color)", color: "", wantErr: false}, + {name: "lowercase rejected", color: "red", wantErr: true}, + {name: "digits rejected", color: "RED1", wantErr: true}, + {name: "punctuation rejected", color: "RED-FOO", wantErr: true}, + } + + for _, tc := range cases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + src := `send [COIN 1] (source = @world \ "` + tc.color + `" destination = @dest)` + parsed := parser.Parse(src) + require.Empty(t, parsed.Errors) + + _, err := machine.RunProgram( + context.Background(), + parsed.Value, + machine.VariablesMap{}, + machine.StaticStore{Balances: machine.Balances{}}, + map[string]struct{}{flags.ExperimentalAssetColors: {}}, + ) + + if tc.wantErr { + require.Error(t, err) + } else { + require.NoError(t, err) + } + }) + } +} + +// Color survives a JSON marshal/unmarshal roundtrip on the Posting type. +func TestPostingJSONRoundtripPreservesColor(t *testing.T) { + t.Parallel() + + src := ` + send [COIN 100] ( + source = @world \ "GRANTS" + destination = @alice + ) + ` + postings, err := runColored(t, src, machine.StaticStore{Balances: machine.Balances{}}) + require.NoError(t, err) + require.Len(t, postings, 1) + require.Equal(t, "GRANTS", postings[0].Color) + + encoded, err := json.Marshal(postings[0]) + require.NoError(t, err) + require.Contains(t, string(encoded), `"color":"GRANTS"`, + "marshaled posting must carry a non-empty color field, got: %s", string(encoded)) + + var decoded machine.Posting + require.NoError(t, json.Unmarshal(encoded, &decoded)) + require.Equal(t, postings[0].Color, decoded.Color) +} + +// Uncolored postings omit the color field — an absent color is the +// uncolored bucket, identical in meaning to color == "". Keeping the JSON +// noise-free for the dominant non-color case. +func TestPostingJSONOmitsEmptyColor(t *testing.T) { + t.Parallel() + + p := machine.Posting{ + Source: "world", + Destination: "dest", + Asset: "COIN", + Amount: big.NewInt(1), + } + encoded, err := json.Marshal(p) + require.NoError(t, err) + require.NotContains(t, string(encoded), `"color"`, + "uncolored postings must not emit a color field, got: %s", string(encoded)) + + // Round-trip: an absent color field decodes back to the uncolored bucket. + var decoded machine.Posting + require.NoError(t, json.Unmarshal(encoded, &decoded)) + require.Equal(t, "", decoded.Color) +} + +// Allocation-style send: one source feeding multiple destinations under a +// color constraint. Every emitted posting must carry the source's color. +func TestColoredAllocationPropagatesToEachLeg(t *testing.T) { + t.Parallel() + + src := ` + send [COIN 100] ( + source = @bank \ "GRANTS" + destination = { + 50% to @alice + remaining to @bob + } + ) + ` + store := machine.StaticStore{Balances: machine.Balances{ + "bank": machine.AccountBalance{ + "COIN": machine.ColorBalance{"GRANTS": big.NewInt(1000)}, + }, + }} + + postings, err := runColored(t, src, store) + require.NoError(t, err) + require.Len(t, postings, 2) + for _, p := range postings { + require.Equal(t, "GRANTS", p.Color, "every leg of the allocation must keep the color") + require.Equal(t, "COIN", p.Asset) + } +} + +// send * (send-all) from a colored bucket drains exactly that bucket and +// emits postings carrying the bucket's color. +func TestColoredSendAllDrainsOnlyTheTargetColor(t *testing.T) { + t.Parallel() + + src := ` + send [COIN *] ( + source = @vault \ "RED" + destination = @dest + ) + ` + store := machine.StaticStore{Balances: machine.Balances{ + "vault": machine.AccountBalance{ + "COIN": machine.ColorBalance{ + "": big.NewInt(1_000_000), + "RED": big.NewInt(42), + "BLUE": big.NewInt(999), + }, + }, + }} + + postings, err := runColored(t, src, store) + require.NoError(t, err) + require.Len(t, postings, 1) + require.Equal(t, "RED", postings[0].Color) + require.Equal(t, big.NewInt(42), postings[0].Amount) +} + +// Sending an empty-color amount must NOT pull from any colored bucket, +// even when the script forms the source from the union of accounts. +func TestUncoloredSourceIgnoresColoredFunds(t *testing.T) { + t.Parallel() + + src := ` + send [COIN 50] ( + source = @vault + destination = @dest + ) + ` + store := machine.StaticStore{Balances: machine.Balances{ + "vault": machine.AccountBalance{ + "COIN": machine.ColorBalance{ + "": big.NewInt(20), // only 20 here — should fail + "RED": big.NewInt(1_000_000), + }, + }, + }} + + _, err := runColored(t, src, store) + require.Error(t, err, "uncolored source must not be able to dip into colored funds") + var missing machine.MissingFundsErr + require.ErrorAs(t, err, &missing) +} + +// Two distinct sources with two distinct colors must each contribute a +// posting bearing their own color (no accidental coalescing). +func TestTwoColoredSourcesYieldTwoColoredPostings(t *testing.T) { + t.Parallel() + + src := ` + send [COIN 60] ( + source = { + @a \ "RED" + @b \ "BLUE" + } + destination = @dest + ) + ` + store := machine.StaticStore{Balances: machine.Balances{ + "a": machine.AccountBalance{"COIN": machine.ColorBalance{"RED": big.NewInt(25)}}, + "b": machine.AccountBalance{"COIN": machine.ColorBalance{"BLUE": big.NewInt(100)}}, + }} + + postings, err := runColored(t, src, store) + require.NoError(t, err) + require.Len(t, postings, 2) + + require.Equal(t, "RED", postings[0].Color) + require.Equal(t, big.NewInt(25), postings[0].Amount) + + require.Equal(t, "BLUE", postings[1].Color) + require.Equal(t, big.NewInt(35), postings[1].Amount) +} + +// Colors play orthogonally with the asset precision suffix (e.g. USD/4) — +// the suffix stays on the asset string, the color rides separately. +func TestColorComposesWithAssetPrecision(t *testing.T) { + t.Parallel() + + src := ` + send [USD/4 10] ( + source = @src \ "COL" allowing unbounded overdraft + destination = @dest + ) + ` + postings, err := runColored(t, src, machine.StaticStore{Balances: machine.Balances{}}) + require.NoError(t, err) + require.Len(t, postings, 1) + require.Equal(t, "USD/4", postings[0].Asset) + require.Equal(t, "COL", postings[0].Color) +} + +// Balances JSON shape — see balances_json.go. +// Three forms are accepted under each (account, asset) entry: +// 1. bare number (uncolored shorthand) +// 2. single value-object {amount, color?} +// 3. array of value-objects (canonical multi-color form) +func TestBalancesJSONShape(t *testing.T) { + t.Parallel() + + src := `{ + "alice": { + "USD/2": [ + { "amount": 100 }, + { "color": "RED", "amount": 50 } + ], + "EUR/2": -42, + "GBP": { "color": "BLUE", "amount": 7 } + } + }` + want := machine.Balances{ + "alice": machine.AccountBalance{ + "USD/2": machine.ColorBalance{"": big.NewInt(100), "RED": big.NewInt(50)}, + "EUR/2": machine.Uncolored(big.NewInt(-42)), + "GBP": machine.ColorBalance{"BLUE": big.NewInt(7)}, + }, + } + + var got machine.Balances + require.NoError(t, json.Unmarshal([]byte(src), &got)) + require.True(t, machine.CompareBalances(want, got), + "unexpected balances: want %v, got %v", want, got) +} + +// Canonical write: uncolored-only collapses to a bare number; multi-color +// emits a deterministic array sorted by color. +func TestBalancesJSONCanonicalWrite(t *testing.T) { + t.Parallel() + + b := machine.Balances{ + "alice": machine.AccountBalance{ + "USD/2": machine.ColorBalance{"": big.NewInt(100), "RED": big.NewInt(50)}, + "EUR/2": machine.Uncolored(big.NewInt(-42)), + }, + } + + encoded, err := json.Marshal(b) + require.NoError(t, err) + + const want = `{"alice":{"EUR/2":-42,"USD/2":[{"amount":100},{"color":"RED","amount":50}]}}` + require.JSONEq(t, want, string(encoded)) +} + +// The old strict-nested shape `{"USD/2": {"": 100}}` is no longer valid: +// the value-object must have an "amount" field. This guards against silent +// acceptance of pre-migration fixtures. +func TestBalancesJSONRejectsLegacyNestedShape(t *testing.T) { + t.Parallel() + + var got machine.Balances + err := json.Unmarshal([]byte(`{"alice": {"USD/2": {"": 100}}}`), &got) + require.Error(t, err, "legacy nested shape must be rejected") +} + +// PrettyPrintPostings must surface Color when any posting carries one — without +// it, two otherwise-identical colored postings would render identically. +func TestPrettyPrintPostingsExposesColor(t *testing.T) { + t.Parallel() + + postings := []machine.Posting{ + {Source: "world", Destination: "alice", Asset: "USD/2", Amount: big.NewInt(100)}, + {Source: "world", Destination: "alice", Asset: "USD/2", Amount: big.NewInt(50), Color: "RED"}, + } + + out := machine.PrettyPrintPostings(postings) + require.Contains(t, out, "Color") + require.Contains(t, out, "RED") +} + +// When no posting is colored, the Color column stays hidden to keep the +// uncolored output unchanged. +func TestPrettyPrintPostingsHidesColorWhenAbsent(t *testing.T) { + t.Parallel() + + postings := []machine.Posting{ + {Source: "world", Destination: "alice", Asset: "USD/2", Amount: big.NewInt(100)}, + } + + out := machine.PrettyPrintPostings(postings) + require.NotContains(t, out, "Color") +} diff --git a/internal/interpreter/funds_queue.go b/internal/interpreter/funds_queue.go index 007c791a..2e4d18c8 100644 --- a/internal/interpreter/funds_queue.go +++ b/internal/interpreter/funds_queue.go @@ -110,18 +110,25 @@ func (s *fundsQueue) Push(senders ...Sender) { } } +// PullAnything is the entry point used by pushReceiver to drain senders +// for a given amount, regardless of color. Each pulled Sender keeps its own +// Color so the receiver-side posting can carry it untouched. func (s *fundsQueue) PullAnything(requiredAmount *big.Int) []Sender { - return s.Pull(requiredAmount, nil) + return s.Pull(requiredAmount) } -func (s *fundsQueue) PullColored(requiredAmount *big.Int, color string) []Sender { - return s.Pull(requiredAmount, &color) -} -func (s *fundsQueue) PullUncolored(requiredAmount *big.Int) []Sender { - return s.PullColored(requiredAmount, "") -} - -func (s *fundsQueue) Pull(requiredAmount *big.Int, color *string) []Sender { +// Pull drains up to requiredAmount from the head of the queue. It does NOT +// verify that the queue holds at least requiredAmount — it simply returns +// whatever is there. The caller is responsible for checking completeness: +// tryTakingExact raises MissingFundsErr when the total sent is below the +// requested amount. +// +// The queue itself is bounded upstream: every pushSender call has already +// been capped by CalculateSafeWithdraw against the source's (asset, color) +// balance, so the queue never holds more than the source can legitimately +// commit. Color is carried by each Sender — Pull preserves it on the way out +// without ever inspecting it. +func (s *fundsQueue) Pull(requiredAmount *big.Int) []Sender { // clone so that we can manipulate this arg requiredAmount = new(big.Int).Set(requiredAmount) @@ -134,16 +141,6 @@ func (s *fundsQueue) Pull(requiredAmount *big.Int, color *string) []Sender { available := s.senders.Head s.senders = s.senders.Tail - if color != nil && available.Color != *color { - out1 := s.Pull(requiredAmount, color) - s.senders = &queue[Sender]{ - Head: available, - Tail: s.senders, - } - out = append(out, out1...) - break - } - switch available.Amount.Cmp(requiredAmount) { case -1: // not enough: out = append(out, available) diff --git a/internal/interpreter/funds_queue_color_test.go b/internal/interpreter/funds_queue_color_test.go new file mode 100644 index 00000000..94c20a76 --- /dev/null +++ b/internal/interpreter/funds_queue_color_test.go @@ -0,0 +1,36 @@ +package interpreter + +// White-box tests for the color-awareness of the internal funds queue. +// These live in `package interpreter` because they reach into unexported +// types (Sender, newFundsQueue). The public color semantics — what the +// numscript ↔ ledger contract actually exposes — are covered in +// color_semantics_test.go (black-box). + +import ( + "math/big" + "testing" + + "github.com/stretchr/testify/require" +) + +// compactTop() relies on (Name, Color) equality: adjacent senders that match +// on both dimensions collapse, those that disagree on either don't. Color +// must NEVER be collapsed across — that would silently merge buckets that +// the source-level segregation is meant to keep apart. +func TestFundsQueueCompactDoesNotMergeAcrossColors(t *testing.T) { + t.Parallel() + + queue := newFundsQueue([]Sender{ + {Name: "a", Color: "RED", Amount: big.NewInt(10)}, + {Name: "a", Color: "RED", Amount: big.NewInt(5)}, + {Name: "a", Color: "BLUE", Amount: big.NewInt(7)}, + }) + + out := queue.PullAnything(big.NewInt(22)) + require.Equal(t, []Sender{ + // (a, RED) entries compact into a single 15 ✓ + {Name: "a", Color: "RED", Amount: big.NewInt(15)}, + // (a, BLUE) stays distinct ✓ + {Name: "a", Color: "BLUE", Amount: big.NewInt(7)}, + }, out) +} diff --git a/internal/interpreter/funds_queue_test.go b/internal/interpreter/funds_queue_test.go index 87672be9..c469a662 100644 --- a/internal/interpreter/funds_queue_test.go +++ b/internal/interpreter/funds_queue_test.go @@ -23,7 +23,7 @@ func TestPush(t *testing.T) { queue := newFundsQueue(nil) queue.Push(Sender{Name: "acc", Amount: big.NewInt(100)}) - out := queue.PullUncolored(big.NewInt(20)) + out := queue.PullAnything(big.NewInt(20)) require.Equal(t, []Sender{ {Name: "acc", Amount: big.NewInt(20)}, }, out) @@ -121,58 +121,6 @@ func TestNoZeroLeftovers(t *testing.T) { }, out) } -func TestReconcileColoredManyDestPerSender(t *testing.T) { - queue := newFundsQueue([]Sender{ - {"src", big.NewInt(10), "X"}, - }) - - out := queue.PullColored(big.NewInt(5), "X") - require.Equal(t, []Sender{ - {Name: "src", Amount: big.NewInt(5), Color: "X"}, - }, out) - - out = queue.PullColored(big.NewInt(5), "X") - require.Equal(t, []Sender{ - {Name: "src", Amount: big.NewInt(5), Color: "X"}, - }, out) - -} - -func TestPullColored(t *testing.T) { - queue := newFundsQueue([]Sender{ - {Name: "s1", Amount: big.NewInt(5)}, - {Name: "s2", Amount: big.NewInt(1), Color: "red"}, - {Name: "s3", Amount: big.NewInt(10)}, - {Name: "s4", Amount: big.NewInt(2), Color: "red"}, - {Name: "s5", Amount: big.NewInt(5)}, - }) - - out := queue.PullColored(big.NewInt(2), "red") - require.Equal(t, []Sender{ - {Name: "s2", Amount: big.NewInt(1), Color: "red"}, - {Name: "s4", Amount: big.NewInt(1), Color: "red"}, - }, out) - - require.Equal(t, []Sender{ - {Name: "s1", Amount: big.NewInt(5)}, - {Name: "s3", Amount: big.NewInt(10)}, - {Name: "s4", Amount: big.NewInt(1), Color: "red"}, - {Name: "s5", Amount: big.NewInt(5)}, - }, queue.PullAll()) -} - -func TestPullColoredComplex(t *testing.T) { - queue := newFundsQueue([]Sender{ - {"s1", big.NewInt(1), "c1"}, - {"s2", big.NewInt(1), "c2"}, - }) - - out := queue.PullColored(big.NewInt(1), "c2") - require.Equal(t, []Sender{ - {Name: "s2", Amount: big.NewInt(1), Color: "c2"}, - }, out) -} - func TestClone(t *testing.T) { fq := newFundsQueue([]Sender{ @@ -190,14 +138,12 @@ func TestClone(t *testing.T) { } func TestCompactFundsAndPush(t *testing.T) { - noCol := "" - queue := newFundsQueue([]Sender{ {Name: "s1", Amount: big.NewInt(2)}, {Name: "s1", Amount: big.NewInt(10)}, }) - queue.Pull(big.NewInt(1), &noCol) + queue.Pull(big.NewInt(1)) queue.Push(Sender{ Name: "pushed", diff --git a/internal/interpreter/interpreter.go b/internal/interpreter/interpreter.go index 94b634af..567a8074 100644 --- a/internal/interpreter/interpreter.go +++ b/internal/interpreter/interpreter.go @@ -16,13 +16,26 @@ import ( type VariablesMap map[string]string -// For each account, list of the needed assets -type BalanceQuery map[string][]string +// AssetColor identifies a (asset, color) pair to query. +// Color is "" when no color is involved. +type AssetColor struct { + Asset string + Color string +} + +// For each account, list of the needed (asset, color) pairs +type BalanceQuery map[string][]AssetColor // For each account, list of the needed keys type MetadataQuery map[string][]string -type AccountBalance = map[string]*big.Int +// ColorBalance maps a color (or "" for the uncolored bucket) to an amount. +type ColorBalance = map[string]*big.Int + +// AccountBalance maps an asset to its per-color balances. +type AccountBalance = map[string]ColorBalance + +// Balances maps an account to its (asset, color) → amount table. type Balances map[string]AccountBalance type AccountMetadata = map[string]string @@ -44,7 +57,7 @@ func (s StaticStore) GetBalances(_ context.Context, q BalanceQuery) (Balances, e } outputBalance := Balances{} - for queriedAccount, queriedCurrencies := range q { + for queriedAccount, queriedItems := range q { outputAccountBalance := AccountBalance{} outputBalance[queriedAccount] = outputAccountBalance @@ -52,27 +65,35 @@ func (s StaticStore) GetBalances(_ context.Context, q BalanceQuery) (Balances, e return AccountBalance{} }) - for _, curr := range queriedCurrencies { - baseAsset, isCatchAll := strings.CutSuffix(curr, "/*") + for _, item := range queriedItems { + baseAsset, isCatchAll := strings.CutSuffix(item.Asset, "/*") if isCatchAll { - - for k, v := range accountBalanceLookup { - matchesAsset := k == baseAsset || strings.HasPrefix(k, baseAsset+"/") + for asset, colorMap := range accountBalanceLookup { + matchesAsset := asset == baseAsset || strings.HasPrefix(asset, baseAsset+"/") if !matchesAsset { continue } - outputAccountBalance[k] = new(big.Int).Set(v) + amt, ok := colorMap[item.Color] + if !ok { + continue + } + out := utils.MapGetOrPutDefault(outputAccountBalance, asset, func() ColorBalance { + return ColorBalance{} + }) + out[item.Color] = new(big.Int).Set(amt) } - } else { + out := utils.MapGetOrPutDefault(outputAccountBalance, item.Asset, func() ColorBalance { + return ColorBalance{} + }) n := new(big.Int) - outputAccountBalance[curr] = n - - if i, ok := accountBalanceLookup[curr]; ok { - n.Set(i) + out[item.Color] = n + if colorMap, ok := accountBalanceLookup[item.Asset]; ok { + if i, ok := colorMap[item.Color]; ok { + n.Set(i) + } } } - } } @@ -98,6 +119,7 @@ type Posting struct { Destination string `json:"destination"` Amount *big.Int `json:"amount"` Asset string `json:"asset"` + Color string `json:"color,omitempty"` } type ExecutionResult struct { @@ -431,7 +453,8 @@ func (st *programState) pushReceiver(name string, monetary *big.Int) { postings := Posting{ Source: sender.Name, Destination: name, - Asset: coloredAsset(st.CurrentAsset, &sender.Color), + Asset: st.CurrentAsset, + Color: sender.Color, Amount: sender.Amount, } @@ -1186,12 +1209,33 @@ func CalculateSafeWithdraw( } func PrettyPrintPostings(postings []Posting) string { + hasColor := false + for _, posting := range postings { + if posting.Color != "" { + hasColor = true + break + } + } + + if !hasColor { + var rows [][]string + for _, posting := range postings { + row := []string{posting.Source, posting.Destination, posting.Asset, posting.Amount.String()} + rows = append(rows, row) + } + return utils.CsvPretty([]string{"Source", "Destination", "Asset", "Amount"}, rows, false) + } + var rows [][]string for _, posting := range postings { - row := []string{posting.Source, posting.Destination, posting.Asset, posting.Amount.String()} + color := posting.Color + if color == "" { + color = "-" + } + row := []string{posting.Source, posting.Destination, posting.Asset, color, posting.Amount.String()} rows = append(rows, row) } - return utils.CsvPretty([]string{"Source", "Destination", "Asset", "Amount"}, rows, false) + return utils.CsvPretty([]string{"Source", "Destination", "Asset", "Color", "Amount"}, rows, false) } func PrettyPrintMeta(meta Metadata) string { diff --git a/internal/interpreter/interpreter_test.go b/internal/interpreter/interpreter_test.go index 651d17b7..b927cd4a 100644 --- a/internal/interpreter/interpreter_test.go +++ b/internal/interpreter/interpreter_test.go @@ -35,25 +35,69 @@ func TestScripts(t *testing.T) { } type TestCase struct { - source string - program *parser.Program - vars map[string]string - meta machine.AccountsMetadata - balances map[string]map[string]*big.Int - expected CaseResult + source string + // program is set by compile() + program *parser.Program + vars map[string]string + meta machine.AccountsMetadata + // balances holds uncolored balances for tests that don't care about color. + // A non-empty color is conveyed via coloredBalances. + balances map[string]map[string]*big.Int + coloredBalances machine.Balances + expected CaseResult } func NewTestCase() TestCase { return TestCase{ - vars: make(map[string]string), - meta: machine.AccountsMetadata{}, - balances: make(map[string]map[string]*big.Int), + vars: make(map[string]string), + meta: machine.AccountsMetadata{}, + balances: make(map[string]map[string]*big.Int), + coloredBalances: machine.Balances{}, expected: CaseResult{ Error: nil, }, } } +// builtBalances merges the uncolored shorthand with any colored balances +// configured on the test case into a single machine.Balances value. +func (c *TestCase) builtBalances() machine.Balances { + out := machine.Balances{} + for acc, assets := range c.balances { + accB, ok := out[acc] + if !ok { + accB = machine.AccountBalance{} + out[acc] = accB + } + for asset, amt := range assets { + colorMap, ok := accB[asset] + if !ok { + colorMap = machine.ColorBalance{} + accB[asset] = colorMap + } + colorMap[""] = amt + } + } + for acc, assets := range c.coloredBalances { + accB, ok := out[acc] + if !ok { + accB = machine.AccountBalance{} + out[acc] = accB + } + for asset, colorMap := range assets { + merged, ok := accB[asset] + if !ok { + merged = machine.ColorBalance{} + accB[asset] = merged + } + for color, amt := range colorMap { + merged[color] = amt + } + } + } + return out +} + // returns a version of the error in which the range is normalized // to golang's default value func removeRange(e machine.InterpreterError) machine.InterpreterError { @@ -104,6 +148,20 @@ func (c *TestCase) setBalance(account string, asset string, amount int64) { c.balances[account][asset] = big.NewInt(amount) } +func (c *TestCase) setColoredBalance(account string, asset string, color string, amount int64) { + accB, ok := c.coloredBalances[account] + if !ok { + accB = machine.AccountBalance{} + c.coloredBalances[account] = accB + } + colorMap, ok := accB[asset] + if !ok { + colorMap = machine.ColorBalance{} + accB[asset] = colorMap + } + colorMap[color] = big.NewInt(amount) +} + func test(t *testing.T, testCase TestCase) { testWithFeatureFlag(t, testCase, "") } @@ -126,8 +184,8 @@ func testWithFeatureFlag(t *testing.T, testCase TestCase, flagName string) { *prog, testCase.vars, machine.StaticStore{ - testCase.balances, - testCase.meta, + Balances: testCase.builtBalances(), + Meta: testCase.meta, }, nil, ) @@ -142,8 +200,8 @@ func testWithFeatureFlag(t *testing.T, testCase TestCase, flagName string) { *prog, testCase.vars, machine.StaticStore{ - testCase.balances, - testCase.meta, + Balances: testCase.builtBalances(), + Meta: testCase.meta, }, featureFlags, ) @@ -161,32 +219,32 @@ func TestStaticStore(t *testing.T) { store := machine.StaticStore{ Balances: machine.Balances{ "a": machine.AccountBalance{ - "USD/2": big.NewInt(10), - "EUR/2": big.NewInt(1), + "USD/2": machine.Uncolored(big.NewInt(10)), + "EUR/2": machine.Uncolored(big.NewInt(1)), }, "b": machine.AccountBalance{ - "USD/2": big.NewInt(10), - "COIN": big.NewInt(11), + "USD/2": machine.Uncolored(big.NewInt(10)), + "COIN": machine.Uncolored(big.NewInt(11)), }, }, } q1, _ := store.GetBalances(context.TODO(), machine.BalanceQuery{ - "a": []string{"USD/2"}, + "a": []machine.AssetColor{{Asset: "USD/2"}}, }) require.Equal(t, machine.Balances{ "a": machine.AccountBalance{ - "USD/2": big.NewInt(10), + "USD/2": machine.Uncolored(big.NewInt(10)), }, }, q1) q2, _ := store.GetBalances(context.TODO(), machine.BalanceQuery{ - "b": []string{"USD/2", "COIN"}, + "b": []machine.AssetColor{{Asset: "USD/2"}, {Asset: "COIN"}}, }) require.Equal(t, machine.Balances{ "b": machine.AccountBalance{ - "USD/2": big.NewInt(10), - "COIN": big.NewInt(11), + "USD/2": machine.Uncolored(big.NewInt(10)), + "COIN": machine.Uncolored(big.NewInt(11)), }, }, q2) }) @@ -195,25 +253,78 @@ func TestStaticStore(t *testing.T) { store := machine.StaticStore{ Balances: machine.Balances{ "a": machine.AccountBalance{ - "USD": big.NewInt(1), - "USD/2": big.NewInt(2), - "USD/3": big.NewInt(3), + "USD": machine.Uncolored(big.NewInt(1)), + "USD/2": machine.Uncolored(big.NewInt(2)), + "USD/3": machine.Uncolored(big.NewInt(3)), }, }, } balances, err := store.GetBalances(context.Background(), machine.BalanceQuery{ - "a": []string{"USD/*"}, + "a": []machine.AssetColor{{Asset: "USD/*"}}, }) require.Nil(t, err) require.Equal(t, machine.Balances{ "a": machine.AccountBalance{ - "USD": big.NewInt(1), - "USD/2": big.NewInt(2), - "USD/3": big.NewInt(3), + "USD": machine.Uncolored(big.NewInt(1)), + "USD/2": machine.Uncolored(big.NewInt(2)), + "USD/3": machine.Uncolored(big.NewInt(3)), + }, + }, balances) + + }) + + t.Run("color dimension is honored", func(t *testing.T) { + store := machine.StaticStore{ + Balances: machine.Balances{ + "a": machine.AccountBalance{ + "USD/2": machine.ColorBalance{ + "": big.NewInt(100), + "GRANTS": big.NewInt(50), + "OPS": big.NewInt(25), + }, + }, + }, + } + + balances, err := store.GetBalances(context.Background(), machine.BalanceQuery{ + "a": []machine.AssetColor{ + {Asset: "USD/2"}, + {Asset: "USD/2", Color: "GRANTS"}, + }, + }) + require.NoError(t, err) + require.Equal(t, machine.Balances{ + "a": machine.AccountBalance{ + "USD/2": machine.ColorBalance{ + "": big.NewInt(100), + "GRANTS": big.NewInt(50), + }, }, }, balances) + }) + + t.Run("color filter is independent across assets in catchall", func(t *testing.T) { + store := machine.StaticStore{ + Balances: machine.Balances{ + "a": machine.AccountBalance{ + "USD": machine.ColorBalance{"": big.NewInt(1), "RED": big.NewInt(10)}, + "USD/2": machine.ColorBalance{"": big.NewInt(2), "RED": big.NewInt(20)}, + "USD/3": machine.ColorBalance{"": big.NewInt(3)}, + }, + }, + } + balances, err := store.GetBalances(context.Background(), machine.BalanceQuery{ + "a": []machine.AssetColor{{Asset: "USD/*", Color: "RED"}}, + }) + require.NoError(t, err) + require.Equal(t, machine.Balances{ + "a": machine.AccountBalance{ + "USD": machine.ColorBalance{"RED": big.NewInt(10)}, + "USD/2": machine.ColorBalance{"RED": big.NewInt(20)}, + }, + }, balances) }) } @@ -727,16 +838,16 @@ func TestTrackBalancesTricky(t *testing.T) { tc.expected = CaseResult{ Postings: []machine.Posting{ { - "world", - "src", - big.NewInt(10), - "GEM", + Source: "world", + Destination: "src", + Amount: big.NewInt(10), + Asset: "GEM", }, { - "src", - "dest", - big.NewInt(15), - "GEM", + Source: "src", + Destination: "dest", + Amount: big.NewInt(15), + Asset: "GEM", }, }, } @@ -1001,7 +1112,7 @@ func TestColorRestrictBalanceWhenMissingFunds(t *testing.T) { tc := NewTestCase() tc.setBalance("acc", "COIN", 100) - tc.setBalance("acc", "COIN_RED", 1) + tc.setColoredBalance("acc", "COIN", "RED", 1) tc.compile(t, script) tc.expected = CaseResult{ diff --git a/internal/interpreter/testdata/numscript-cookbook/experimental/mixed-source-prefer-single-source.num.specs.json b/internal/interpreter/testdata/numscript-cookbook/experimental/mixed-source-prefer-single-source.num.specs.json index c54424ab..f8e22097 100644 --- a/internal/interpreter/testdata/numscript-cookbook/experimental/mixed-source-prefer-single-source.num.specs.json +++ b/internal/interpreter/testdata/numscript-cookbook/experimental/mixed-source-prefer-single-source.num.specs.json @@ -3,7 +3,9 @@ "variables": { "amt": "10" }, - "featureFlags": ["experimental-oneof"], + "featureFlags": [ + "experimental-oneof" + ], "testCases": [ { "it": "sends from the first source when there is enough balance", diff --git a/internal/interpreter/testdata/numscript-cookbook/experimental/top-up-many.num.specs.json b/internal/interpreter/testdata/numscript-cookbook/experimental/top-up-many.num.specs.json index 2c1cbc95..7bdbcf61 100644 --- a/internal/interpreter/testdata/numscript-cookbook/experimental/top-up-many.num.specs.json +++ b/internal/interpreter/testdata/numscript-cookbook/experimental/top-up-many.num.specs.json @@ -11,14 +11,25 @@ "testCases": [ { "it": "should be a noop when all balances are >= 0", - "balances": { "alice": { "EUR": 100 }, "bob": { "EUR": 200 } }, + "balances": { + "alice": { + "EUR": 100 + }, + "bob": { + "EUR": 200 + } + }, "expect.postings": [] }, { "it": "should prioritize alice when both have missing funds", "balances": { - "alice": { "EUR": -120 }, - "bob": { "EUR": -120 } + "alice": { + "EUR": -120 + }, + "bob": { + "EUR": -120 + } }, "expect.postings": [ { @@ -32,8 +43,12 @@ { "it": "doesn't send funds to alice if there aren't enough funds for the account to be topped-up", "balances": { - "alice": { "EUR": -80 }, - "bob": { "EUR": -120 } + "alice": { + "EUR": -80 + }, + "bob": { + "EUR": -120 + } }, "expect.postings": [ { @@ -47,8 +62,12 @@ { "it": "funds are kept if there are spare funds", "balances": { - "alice": { "EUR": -10 }, - "bob": { "EUR": -20 } + "alice": { + "EUR": -10 + }, + "bob": { + "EUR": -20 + } }, "expect.postings": [] } diff --git a/internal/interpreter/testdata/numscript-cookbook/experimental/top-up.num.specs.json b/internal/interpreter/testdata/numscript-cookbook/experimental/top-up.num.specs.json index 05fcf681..70e60b1d 100644 --- a/internal/interpreter/testdata/numscript-cookbook/experimental/top-up.num.specs.json +++ b/internal/interpreter/testdata/numscript-cookbook/experimental/top-up.num.specs.json @@ -1,22 +1,38 @@ { "$schema": "https://raw.githubusercontent.com/formancehq/numscript/main/specs.schema.json", - "featureFlags": ["experimental-overdraft-function"], + "featureFlags": [ + "experimental-overdraft-function" + ], "testCases": [ { "it": "should not emit postings", - "balances": { "alice": { "EUR": 100 } }, + "balances": { + "alice": { + "EUR": 100 + } + }, "expect.postings": [] }, { "it": "should send the missing amount to an overdraft account", - "balances": { "alice": { "EUR": -100 } }, + "balances": { + "alice": { + "EUR": -100 + } + }, "expect.endBalances": { - "alice": { "EUR": 0 }, - "world": { "EUR": -100 } + "alice": { + "EUR": 0 + }, + "world": { + "EUR": -100 + } }, "expect.movements": { "world": { - "alice": { "EUR": 100 } + "alice": { + "EUR": 100 + } } }, "expect.postings": [ diff --git a/internal/interpreter/testdata/numscript-cookbook/experimental/transfer-example.num.specs.json b/internal/interpreter/testdata/numscript-cookbook/experimental/transfer-example.num.specs.json index c0e92c9a..9502a895 100644 --- a/internal/interpreter/testdata/numscript-cookbook/experimental/transfer-example.num.specs.json +++ b/internal/interpreter/testdata/numscript-cookbook/experimental/transfer-example.num.specs.json @@ -13,8 +13,12 @@ { "it": "should authorize transfer if both wallet and bank accounts display enough balance", "balances": { - "wallet": { "EUR": 100 }, - "bank_account": { "EUR": -100 } + "wallet": { + "EUR": 100 + }, + "bank_account": { + "EUR": -100 + } }, "expect.postings": [ { @@ -28,16 +32,24 @@ { "it": "should not authorize transfer if wallet does not display enough balance and bank account does", "balances": { - "wallet": { "EUR": 50 }, - "bank_account": { "EUR": -100 } + "wallet": { + "EUR": 50 + }, + "bank_account": { + "EUR": -100 + } }, "expect.error.missingFunds": true }, { "it": "should not authorize transfer if bank account does not display enough balance and wallet does", "balances": { - "wallet": { "EUR": 100 }, - "bank_account": { "EUR": -50 } + "wallet": { + "EUR": 100 + }, + "bank_account": { + "EUR": -50 + } }, "expect.error.missingFunds": true } diff --git a/internal/interpreter/testdata/numscript-cookbook/send-max.num.specs.json b/internal/interpreter/testdata/numscript-cookbook/send-max.num.specs.json index d2e32539..92068884 100644 --- a/internal/interpreter/testdata/numscript-cookbook/send-max.num.specs.json +++ b/internal/interpreter/testdata/numscript-cookbook/send-max.num.specs.json @@ -4,7 +4,9 @@ { "it": "is capped to USD100 when balance is higher", "balances": { - "src": { "USD": 999 } + "src": { + "USD": 999 + } }, "expect.postings": [ { @@ -18,7 +20,9 @@ { "it": "allows sending less than the cap when balance is not enough", "balances": { - "src": { "USD": 42 } + "src": { + "USD": 42 + } }, "expect.postings": [ { diff --git a/internal/interpreter/testdata/script-tests/balance-not-found.num.specs.json b/internal/interpreter/testdata/script-tests/balance-not-found.num.specs.json index c535912d..0ef4b1a7 100644 --- a/internal/interpreter/testdata/script-tests/balance-not-found.num.specs.json +++ b/internal/interpreter/testdata/script-tests/balance-not-found.num.specs.json @@ -8,4 +8,4 @@ "expect.postings": [] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/bigint-literal.num.specs.json b/internal/interpreter/testdata/script-tests/bigint-literal.num.specs.json index a33b882a..9bfb1cb8 100644 --- a/internal/interpreter/testdata/script-tests/bigint-literal.num.specs.json +++ b/internal/interpreter/testdata/script-tests/bigint-literal.num.specs.json @@ -13,4 +13,4 @@ ] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/empty-postings.num.specs.json b/internal/interpreter/testdata/script-tests/empty-postings.num.specs.json index e087b2e8..e1a6dc86 100644 --- a/internal/interpreter/testdata/script-tests/empty-postings.num.specs.json +++ b/internal/interpreter/testdata/script-tests/empty-postings.num.specs.json @@ -10,4 +10,4 @@ "expect.postings": [] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/experimental/account-interpolation/account-interp.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/account-interpolation/account-interp.num.specs.json index 97effe85..fe8165f8 100644 --- a/internal/interpreter/testdata/script-tests/experimental/account-interpolation/account-interp.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/account-interpolation/account-interp.num.specs.json @@ -1,5 +1,7 @@ { - "featureFlags": ["experimental-account-interpolation"], + "featureFlags": [ + "experimental-account-interpolation" + ], "testCases": [ { "it": "-", diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-inorder-send-all.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-inorder-send-all.num.specs.json index 4efe1917..78452622 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-inorder-send-all.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-inorder-send-all.num.specs.json @@ -7,9 +7,19 @@ "it": "-", "balances": { "src": { - "COIN": 100, - "COIN_BLUE": 30, - "COIN_RED": 20 + "COIN": [ + { + "amount": 100 + }, + { + "color": "BLUE", + "amount": 30 + }, + { + "color": "RED", + "amount": 20 + } + ] } }, "expect.postings": [ @@ -17,13 +27,15 @@ "source": "src", "destination": "dest", "amount": 20, - "asset": "COIN_RED" + "asset": "COIN", + "color": "RED" }, { "source": "src", "destination": "dest", "amount": 30, - "asset": "COIN_BLUE" + "asset": "COIN", + "color": "BLUE" }, { "source": "src", diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-inorder.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-inorder.num.specs.json index 2b306b8f..0470c4f2 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-inorder.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-inorder.num.specs.json @@ -7,9 +7,19 @@ "it": "-", "balances": { "src": { - "COIN": 100, - "COIN_BLUE": 30, - "COIN_RED": 20 + "COIN": [ + { + "amount": 100 + }, + { + "color": "BLUE", + "amount": 30 + }, + { + "color": "RED", + "amount": 20 + } + ] } }, "expect.postings": [ @@ -17,13 +27,15 @@ "source": "src", "destination": "dest", "amount": 20, - "asset": "COIN_RED" + "asset": "COIN", + "color": "RED" }, { "source": "src", "destination": "dest", "amount": 30, - "asset": "COIN_BLUE" + "asset": "COIN", + "color": "BLUE" }, { "source": "src", diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-restrict-balance-when-missing-funds.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-restrict-balance-when-missing-funds.num.specs.json index f71564aa..b761ac6f 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-restrict-balance-when-missing-funds.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-restrict-balance-when-missing-funds.num.specs.json @@ -1,12 +1,21 @@ { - "featureFlags": ["experimental-asset-colors"], + "featureFlags": [ + "experimental-asset-colors" + ], "testCases": [ { "it": "-", "balances": { "acc": { - "COIN": 100, - "COIN_RED": 1 + "COIN": [ + { + "amount": 100 + }, + { + "color": "RED", + "amount": 1 + } + ] } }, "expect.error.missingFunds": true diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-restrict-balance.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-restrict-balance.num.specs.json index a4d32470..b3c84dc6 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-restrict-balance.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-restrict-balance.num.specs.json @@ -7,8 +7,15 @@ "it": "-", "balances": { "acc": { - "COIN": 1, - "COIN_RED": 100 + "COIN": [ + { + "amount": 1 + }, + { + "color": "RED", + "amount": 100 + } + ] } }, "expect.postings": [ @@ -16,7 +23,8 @@ "source": "acc", "destination": "dest", "amount": 20, - "asset": "COIN_RED" + "asset": "COIN", + "color": "RED" } ] } diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-restriction-in-send-all.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-restriction-in-send-all.num.specs.json index 8bde6bea..3bc02a31 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-restriction-in-send-all.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-restriction-in-send-all.num.specs.json @@ -7,7 +7,10 @@ "it": "-", "balances": { "src": { - "COIN_RED": 42 + "COIN": { + "color": "RED", + "amount": 42 + } } }, "expect.postings": [ @@ -15,7 +18,8 @@ "source": "src", "destination": "dest", "amount": 42, - "asset": "COIN_RED" + "asset": "COIN", + "color": "RED" } ] } diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-send-overdrat.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-send-overdrat.num.specs.json index 5cfde80f..1c824d8e 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-send-overdrat.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-send-overdrat.num.specs.json @@ -10,7 +10,8 @@ "source": "acc", "destination": "dest", "amount": 100, - "asset": "COIN_RED" + "asset": "COIN", + "color": "RED" } ] } diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-send.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-send.num.specs.json index 4288aed7..31d2e41b 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-send.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-send.num.specs.json @@ -10,7 +10,8 @@ "source": "world", "destination": "dest", "amount": 100, - "asset": "COIN_RED" + "asset": "COIN", + "color": "RED" } ] } diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-with-asset-precision.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-with-asset-precision.num.specs.json index f9b3821e..3bbe0218 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-with-asset-precision.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-colors/color-with-asset-precision.num.specs.json @@ -10,7 +10,8 @@ "source": "src", "destination": "dest", "amount": 10, - "asset": "USD_COL/4" + "asset": "USD/4", + "color": "COL" } ] } diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-colors/no-double-spending-in-colored-send-all.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-colors/no-double-spending-in-colored-send-all.num.specs.json index a1b7d2f9..af5fb491 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-colors/no-double-spending-in-colored-send-all.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-colors/no-double-spending-in-colored-send-all.num.specs.json @@ -7,8 +7,15 @@ "it": "-", "balances": { "src": { - "COIN": 100, - "COIN_X": 20 + "COIN": [ + { + "amount": 100 + }, + { + "color": "X", + "amount": 20 + } + ] } }, "expect.postings": [ @@ -16,7 +23,8 @@ "source": "src", "destination": "dest", "amount": 20, - "asset": "COIN_X" + "asset": "COIN", + "color": "X" }, { "source": "src", diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-colors/no-double-spending-in-colored-send.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-colors/no-double-spending-in-colored-send.num.specs.json index da4b831a..709e60f4 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-colors/no-double-spending-in-colored-send.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-colors/no-double-spending-in-colored-send.num.specs.json @@ -7,8 +7,15 @@ "it": "-", "balances": { "src": { - "COIN": 99999, - "COIN_X": 20 + "COIN": [ + { + "amount": 99999 + }, + { + "color": "X", + "amount": 20 + } + ] } }, "expect.postings": [ @@ -16,7 +23,8 @@ "source": "src", "destination": "dest", "amount": 20, - "asset": "COIN_X" + "asset": "COIN", + "color": "X" }, { "source": "src", diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-scaling/no-solution.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-scaling/no-solution.num.specs.json index 0d5137db..3ca0b4c9 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-scaling/no-solution.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-scaling/no-solution.num.specs.json @@ -74,4 +74,4 @@ "expect.error.missingFunds": true } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-all-allotment.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-all-allotment.num.specs.json index ce35e9ea..b0ea06b4 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-all-allotment.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-all-allotment.num.specs.json @@ -44,4 +44,4 @@ ] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-allotment.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-allotment.num.specs.json index bb464f4c..da48ec98 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-allotment.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-allotment.num.specs.json @@ -100,4 +100,4 @@ ] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-kept.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-kept.num.specs.json index 102fc196..fea0689f 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-kept.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-kept.num.specs.json @@ -28,4 +28,4 @@ ] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-send-all.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-send-all.num.specs.json index 75c32eea..694a0645 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-send-all.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-send-all.num.specs.json @@ -69,4 +69,4 @@ ] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-with-oneof.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-with-oneof.num.specs.json index 68bb6697..654b91d0 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-with-oneof.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling-with-oneof.num.specs.json @@ -50,4 +50,4 @@ ] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling.num.specs.json index 1799fe48..ecec370f 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-scaling/scaling.num.specs.json @@ -102,4 +102,4 @@ "expect.error.missingFunds": true } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/experimental/asset-scaling/update-swap-account-balance.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/asset-scaling/update-swap-account-balance.num.specs.json index 52c0e5bf..700f49e5 100644 --- a/internal/interpreter/testdata/script-tests/experimental/asset-scaling/update-swap-account-balance.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/asset-scaling/update-swap-account-balance.num.specs.json @@ -51,4 +51,4 @@ ] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/experimental/get-amount-function/get-amount-function.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/get-amount-function/get-amount-function.num.specs.json index d52e8c6e..32d21063 100644 --- a/internal/interpreter/testdata/script-tests/experimental/get-amount-function/get-amount-function.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/get-amount-function/get-amount-function.num.specs.json @@ -1,5 +1,7 @@ { - "featureFlags": ["experimental-get-amount-function"], + "featureFlags": [ + "experimental-get-amount-function" + ], "testCases": [ { "it": "-", diff --git a/internal/interpreter/testdata/script-tests/experimental/get-asset-function/get-asset-function.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/get-asset-function/get-asset-function.num.specs.json index 3b3997b7..3e83b43e 100644 --- a/internal/interpreter/testdata/script-tests/experimental/get-asset-function/get-asset-function.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/get-asset-function/get-asset-function.num.specs.json @@ -1,5 +1,7 @@ { - "featureFlags": ["experimental-get-asset-function"], + "featureFlags": [ + "experimental-get-asset-function" + ], "testCases": [ { "it": "-", diff --git a/internal/interpreter/testdata/script-tests/experimental/mid-script-function-call/expr-in-var-origin.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/mid-script-function-call/expr-in-var-origin.num.specs.json index 84cb2b44..d0853e40 100644 --- a/internal/interpreter/testdata/script-tests/experimental/mid-script-function-call/expr-in-var-origin.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/mid-script-function-call/expr-in-var-origin.num.specs.json @@ -8,4 +8,4 @@ "expect.postings": [] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/experimental/oneof/oneof-all-failing.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/oneof/oneof-all-failing.num.specs.json index a506b5eb..18f41f81 100644 --- a/internal/interpreter/testdata/script-tests/experimental/oneof/oneof-all-failing.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/oneof/oneof-all-failing.num.specs.json @@ -1,5 +1,7 @@ { - "featureFlags": ["experimental-oneof"], + "featureFlags": [ + "experimental-oneof" + ], "testCases": [ { "it": "-", diff --git a/internal/interpreter/testdata/script-tests/experimental/overdraft-function/overdraft-function-when-positive.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/overdraft-function/overdraft-function-when-positive.num.specs.json index d3093d3d..d5a1d0c4 100644 --- a/internal/interpreter/testdata/script-tests/experimental/overdraft-function/overdraft-function-when-positive.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/overdraft-function/overdraft-function-when-positive.num.specs.json @@ -13,4 +13,4 @@ "expect.postings": [] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/experimental/overdraft-function/overdraft-function-when-zero.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/overdraft-function/overdraft-function-when-zero.num.specs.json index 201383c1..7db029e3 100644 --- a/internal/interpreter/testdata/script-tests/experimental/overdraft-function/overdraft-function-when-zero.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/overdraft-function/overdraft-function-when-zero.num.specs.json @@ -11,4 +11,4 @@ "expect.postings": [] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/experimental/overdraft-function/reach-zero.num.specs.json b/internal/interpreter/testdata/script-tests/experimental/overdraft-function/reach-zero.num.specs.json index 7cc09268..30d34d68 100644 --- a/internal/interpreter/testdata/script-tests/experimental/overdraft-function/reach-zero.num.specs.json +++ b/internal/interpreter/testdata/script-tests/experimental/overdraft-function/reach-zero.num.specs.json @@ -3,11 +3,15 @@ "variables": { "amt": "100" }, - "featureFlags": ["experimental-overdraft-function"], + "featureFlags": [ + "experimental-overdraft-function" + ], "testCases": [ { "balances": { - "invoice:001": { "USD/2": -999 } + "invoice:001": { + "USD/2": -999 + } }, "it": "only sends to the first one if there's not enough funds to top-up its balance", "expect.postings": [ @@ -21,8 +25,12 @@ }, { "balances": { - "invoice:001": { "USD/2": -99 }, - "invoice:002": { "USD/2": -999 } + "invoice:001": { + "USD/2": -99 + }, + "invoice:002": { + "USD/2": -999 + } }, "it": "it sends to the second source after topping-up the first one", "expect.postings": [ @@ -42,8 +50,12 @@ }, { "balances": { - "invoice:001": { "USD/2": -2 }, - "invoice:002": { "USD/2": -3 } + "invoice:001": { + "USD/2": -2 + }, + "invoice:002": { + "USD/2": -3 + } }, "it": "it keeps the spare amount", "expect.postings": [ diff --git a/internal/interpreter/testdata/script-tests/feature-flag-syntax.num.specs.json b/internal/interpreter/testdata/script-tests/feature-flag-syntax.num.specs.json index db6f647b..6b947712 100644 --- a/internal/interpreter/testdata/script-tests/feature-flag-syntax.num.specs.json +++ b/internal/interpreter/testdata/script-tests/feature-flag-syntax.num.specs.json @@ -18,4 +18,4 @@ ] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/floating-perc.num.specs.json b/internal/interpreter/testdata/script-tests/floating-perc.num.specs.json index 185ccf3d..67eba0dd 100644 --- a/internal/interpreter/testdata/script-tests/floating-perc.num.specs.json +++ b/internal/interpreter/testdata/script-tests/floating-perc.num.specs.json @@ -13,4 +13,4 @@ ] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/floating-perc2.num.specs.json b/internal/interpreter/testdata/script-tests/floating-perc2.num.specs.json index 4c512b52..6aa8f68e 100644 --- a/internal/interpreter/testdata/script-tests/floating-perc2.num.specs.json +++ b/internal/interpreter/testdata/script-tests/floating-perc2.num.specs.json @@ -13,4 +13,4 @@ ] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/minus-infix-monetary.num.specs.json b/internal/interpreter/testdata/script-tests/minus-infix-monetary.num.specs.json index 0727d9b2..50c15bdb 100644 --- a/internal/interpreter/testdata/script-tests/minus-infix-monetary.num.specs.json +++ b/internal/interpreter/testdata/script-tests/minus-infix-monetary.num.specs.json @@ -13,4 +13,4 @@ ] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/minus-infix-number.num.specs.json b/internal/interpreter/testdata/script-tests/minus-infix-number.num.specs.json index 336bb433..492863a4 100644 --- a/internal/interpreter/testdata/script-tests/minus-infix-number.num.specs.json +++ b/internal/interpreter/testdata/script-tests/minus-infix-number.num.specs.json @@ -13,4 +13,4 @@ ] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/minus-prefix-number.num.specs.json b/internal/interpreter/testdata/script-tests/minus-prefix-number.num.specs.json index 68dac993..0e3abba8 100644 --- a/internal/interpreter/testdata/script-tests/minus-prefix-number.num.specs.json +++ b/internal/interpreter/testdata/script-tests/minus-prefix-number.num.specs.json @@ -13,4 +13,4 @@ ] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/neg-max-dest.num.specs.json b/internal/interpreter/testdata/script-tests/neg-max-dest.num.specs.json index 51120ad9..656e8d3f 100644 --- a/internal/interpreter/testdata/script-tests/neg-max-dest.num.specs.json +++ b/internal/interpreter/testdata/script-tests/neg-max-dest.num.specs.json @@ -3,7 +3,9 @@ { "it": "-", "balances": { - "memo:main": { "EUR/2": 99999999 } + "memo:main": { + "EUR/2": 99999999 + } }, "variables": {}, "metadata": {}, diff --git a/internal/interpreter/testdata/script-tests/negative-max-send-all.num.specs.json b/internal/interpreter/testdata/script-tests/negative-max-send-all.num.specs.json index 6f5984b0..f50a6f05 100644 --- a/internal/interpreter/testdata/script-tests/negative-max-send-all.num.specs.json +++ b/internal/interpreter/testdata/script-tests/negative-max-send-all.num.specs.json @@ -10,4 +10,4 @@ "expect.postings": [] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/send-all-when-negative.num.specs.json b/internal/interpreter/testdata/script-tests/send-all-when-negative.num.specs.json index 98bdf829..298ea1ac 100644 --- a/internal/interpreter/testdata/script-tests/send-all-when-negative.num.specs.json +++ b/internal/interpreter/testdata/script-tests/send-all-when-negative.num.specs.json @@ -10,4 +10,4 @@ "expect.postings": [] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/send-allt-max-when-no-amount.num.specs.json b/internal/interpreter/testdata/script-tests/send-allt-max-when-no-amount.num.specs.json index 98e41836..50a034c2 100644 --- a/internal/interpreter/testdata/script-tests/send-allt-max-when-no-amount.num.specs.json +++ b/internal/interpreter/testdata/script-tests/send-allt-max-when-no-amount.num.specs.json @@ -11,4 +11,4 @@ "expect.postings": [] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/send-zero.num.specs.json b/internal/interpreter/testdata/script-tests/send-zero.num.specs.json index 5d6aaaa7..78c8bd3d 100644 --- a/internal/interpreter/testdata/script-tests/send-zero.num.specs.json +++ b/internal/interpreter/testdata/script-tests/send-zero.num.specs.json @@ -8,4 +8,4 @@ "expect.postings": [] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/zero-postings-explicit-allotment.num.specs.json b/internal/interpreter/testdata/script-tests/zero-postings-explicit-allotment.num.specs.json index 14cd606a..46ef4203 100644 --- a/internal/interpreter/testdata/script-tests/zero-postings-explicit-allotment.num.specs.json +++ b/internal/interpreter/testdata/script-tests/zero-postings-explicit-allotment.num.specs.json @@ -9,4 +9,4 @@ "expect.postings": [] } ] -} \ No newline at end of file +} diff --git a/internal/interpreter/testdata/script-tests/zero-postings-explicit-inorder.num.specs.json b/internal/interpreter/testdata/script-tests/zero-postings-explicit-inorder.num.specs.json index 23428b37..32812803 100644 --- a/internal/interpreter/testdata/script-tests/zero-postings-explicit-inorder.num.specs.json +++ b/internal/interpreter/testdata/script-tests/zero-postings-explicit-inorder.num.specs.json @@ -10,4 +10,4 @@ "expect.postings": [] } ] -} \ No newline at end of file +} diff --git a/internal/mcp_impl/handlers.go b/internal/mcp_impl/handlers.go index 93139eb9..be5f1c21 100644 --- a/internal/mcp_impl/handlers.go +++ b/internal/mcp_impl/handlers.go @@ -2,8 +2,8 @@ package mcp_impl import ( "context" + "encoding/json" "fmt" - "math/big" "strings" "github.com/formancehq/numscript/internal/analysis" @@ -13,34 +13,28 @@ import ( "github.com/mark3labs/mcp-go/server" ) +// parseBalancesJson routes the caller-supplied balances payload through the +// canonical interpreter.Balances unmarshaller, so the MCP `evaluate` tool +// accepts exactly the same shapes documented elsewhere — bare number for the +// uncolored shorthand, single value-object, or array of value-objects. +// +// The MCP framework hands us an already-decoded `any`; re-marshalling round +// trips through JSON, but big-integer amounts stay safe because the +// re-encoded payload feeds into Balances.UnmarshalJSON which decodes each +// amount directly into a *big.Int. func parseBalancesJson(balancesRaw any) (interpreter.Balances, *mcp.CallToolResult) { - balances, ok := balancesRaw.(map[string]any) - if !ok { - return interpreter.Balances{}, mcp.NewToolResultError(fmt.Sprintf("Expected an object as balances, got: <%#v>", balancesRaw)) + if balancesRaw == nil { + return interpreter.Balances{}, nil } - - iBalances := interpreter.Balances{} - for account, assetsRaw := range balances { - if iBalances[account] == nil { - iBalances[account] = interpreter.AccountBalance{} - } - - assets, ok := assetsRaw.(map[string]any) - if !ok { - return interpreter.Balances{}, mcp.NewToolResultError(fmt.Sprintf("Expected nested object for account %v", account)) - } - - for asset, amountRaw := range assets { - amount, ok := amountRaw.(float64) - if !ok { - return interpreter.Balances{}, mcp.NewToolResultError(fmt.Sprintf("Expected float for amount: %v", amountRaw)) - } - - n, _ := big.NewFloat(amount).Int(new(big.Int)) - iBalances[account][asset] = n - } + encoded, err := json.Marshal(balancesRaw) + if err != nil { + return interpreter.Balances{}, mcp.NewToolResultError(fmt.Sprintf("Could not re-encode balances: %v", err)) + } + var balances interpreter.Balances + if err := json.Unmarshal(encoded, &balances); err != nil { + return interpreter.Balances{}, mcp.NewToolResultError(fmt.Sprintf("Invalid balances payload: %v", err)) } - return iBalances, nil + return balances, nil } func parseVarsJson(varsRaw any) (map[string]string, *mcp.CallToolResult) { @@ -75,8 +69,19 @@ func addEvalTool(s *server.MCPServer) { ), mcp.WithObject("balances", mcp.Required(), - mcp.Description(`The accounts' balances. A nested map from the account name, to the asset, to its integer amount. - For example: { "alice": { "USD/2": 100, "EUR/2": -42 }, "bob": { "BTC": 1 } } + mcp.Description(`The accounts' balances. A nested map from the account name, to the asset, to the held amount. + + Each per-asset entry accepts three forms: + - a bare integer for the uncolored bucket (shorthand) + - a single value-object: { "amount": N } or { "color": "RED", "amount": N } + - an array of value-objects when several colors coexist on the same asset + + Examples: + { "alice": { "USD/2": 100, "EUR/2": -42 }, "bob": { "BTC": 1 } } + { "alice": { "USD/2": [{ "amount": 100 }, { "color": "RED", "amount": 50 }] } } + + Color is a first-class dimension on the emitted postings (see Posting.color); + the empty/missing color is its own bucket, distinct from any non-empty one. `), ), mcp.WithObject("vars", diff --git a/internal/specs_format/__snapshots__/runner_test.snap b/internal/specs_format/__snapshots__/runner_test.snap index 8eab7c2a..2836fed4 100755 --- a/internal/specs_format/__snapshots__/runner_test.snap +++ b/internal/specs_format/__snapshots__/runner_test.snap @@ -53,8 +53,8 @@ GOT:  GIVEN: -| Account | Asset | Balance | -| alice | USD/2 | 9999 | +| Account | Asset | Color | Balance | +| alice | USD/2 | | 9999 |  GOT: @@ -118,7 +118,7 @@ Error: example.num.specs.json  Error: example.num.specs.json -json: cannot unmarshal number into Go struct field Specs.balances of type interpreter.Balances +json: cannot unmarshal number into Go struct field Specs.balances of type map[string]map[string]json.RawMessage --- diff --git a/internal/specs_format/index.go b/internal/specs_format/index.go index cb9b48d3..c42ffe98 100644 --- a/internal/specs_format/index.go +++ b/internal/specs_format/index.go @@ -299,17 +299,25 @@ func getMovements(postings []interpreter.Posting) Movements { return m } +func fetchOrInsertBalance(b interpreter.Balances, account, asset, color string) *big.Int { + accBalance := utils.MapGetOrPutDefault(b, account, func() interpreter.AccountBalance { + return interpreter.AccountBalance{} + }) + colorMap := utils.MapGetOrPutDefault(accBalance, asset, func() interpreter.ColorBalance { + return interpreter.ColorBalance{} + }) + return utils.MapGetOrPutDefault(colorMap, color, func() *big.Int { + return new(big.Int) + }) +} + func getBalances(postings []interpreter.Posting, initialBalances interpreter.Balances) interpreter.Balances { balances := initialBalances.DeepClone() for _, posting := range postings { - sourceBalance := utils.NestedMapGetOrPutDefault(balances, posting.Source, posting.Asset, func() *big.Int { - return new(big.Int) - }) + sourceBalance := fetchOrInsertBalance(balances, posting.Source, posting.Asset, posting.Color) sourceBalance.Sub(sourceBalance, posting.Amount) - destinationBalance := utils.NestedMapGetOrPutDefault(balances, posting.Destination, posting.Asset, func() *big.Int { - return new(big.Int) - }) + destinationBalance := fetchOrInsertBalance(balances, posting.Destination, posting.Asset, posting.Color) destinationBalance.Add(destinationBalance, posting.Amount) } diff --git a/internal/specs_format/parse_test.go b/internal/specs_format/parse_test.go index 71a91f6f..ce7f2e14 100644 --- a/internal/specs_format/parse_test.go +++ b/internal/specs_format/parse_test.go @@ -47,7 +47,7 @@ func TestParseSpecs(t *testing.T) { require.Equal(t, specs_format.Specs{ Balances: interpreter.Balances{ "alice": { - "EUR": big.NewInt(200), + "EUR": interpreter.Uncolored(big.NewInt(200)), }, }, Vars: interpreter.VariablesMap{ @@ -58,7 +58,7 @@ func TestParseSpecs(t *testing.T) { It: "d1", Balances: interpreter.Balances{ "bob": { - "EUR": big.NewInt(42), + "EUR": interpreter.Uncolored(big.NewInt(42)), }, }, ExpectPostings: []interpreter.Posting{ diff --git a/internal/specs_format/run_test.go b/internal/specs_format/run_test.go index dc7be8bd..0bb505e5 100644 --- a/internal/specs_format/run_test.go +++ b/internal/specs_format/run_test.go @@ -58,7 +58,7 @@ func TestRunSpecsSimple(t *testing.T) { }, Balances: interpreter.Balances{ "src": interpreter.AccountBalance{ - "USD": big.NewInt(9999), + "USD": interpreter.Uncolored(big.NewInt(9999)), }, }, Meta: interpreter.AccountsMetadata{}, @@ -119,11 +119,11 @@ func TestRunSpecsMergeOuter(t *testing.T) { Meta: interpreter.AccountsMetadata{}, Balances: interpreter.Balances{ "src": interpreter.AccountBalance{ - "USD": big.NewInt(10), - "EUR": big.NewInt(2), + "USD": interpreter.Uncolored(big.NewInt(10)), + "EUR": interpreter.Uncolored(big.NewInt(2)), }, "dest": interpreter.AccountBalance{ - "USD": big.NewInt(1), + "USD": interpreter.Uncolored(big.NewInt(1)), }, }, FailedAssertions: nil, @@ -175,7 +175,7 @@ func TestRunWithMissingBalance(t *testing.T) { }, Balances: interpreter.Balances{ "src": interpreter.AccountBalance{ - "USD": big.NewInt(1), + "USD": interpreter.Uncolored(big.NewInt(1)), }, }, Meta: interpreter.AccountsMetadata{}, @@ -229,7 +229,7 @@ func TestRunWithMissingBalanceWhenExpectedPostings(t *testing.T) { }, Balances: interpreter.Balances{ "src": interpreter.AccountBalance{ - "USD": big.NewInt(1), + "USD": interpreter.Uncolored(big.NewInt(1)), }, }, Meta: interpreter.AccountsMetadata{}, @@ -281,7 +281,7 @@ func TestNullPostingsIsNoop(t *testing.T) { }, Balances: interpreter.Balances{ "src": interpreter.AccountBalance{ - "USD": big.NewInt(1), + "USD": interpreter.Uncolored(big.NewInt(1)), }, }, Meta: interpreter.AccountsMetadata{}, @@ -422,7 +422,7 @@ func TestFocus(t *testing.T) { }, Balances: interpreter.Balances{ "src": interpreter.AccountBalance{ - "USD": big.NewInt(9999), + "USD": interpreter.Uncolored(big.NewInt(9999)), }, }, Meta: interpreter.AccountsMetadata{}, diff --git a/numscript.go b/numscript.go index e7fa0437..a74b4a34 100644 --- a/numscript.go +++ b/numscript.go @@ -52,10 +52,13 @@ type ( VariablesMap = interpreter.VariablesMap Posting = interpreter.Posting ExecutionResult = interpreter.ExecutionResult - // For each account, list of the needed assets - BalanceQuery = interpreter.BalanceQuery + // For each account, list of the needed (asset, color) pairs + BalanceQuery = interpreter.BalanceQuery + // AssetColor identifies a (asset, color) pair to query. + AssetColor = interpreter.AssetColor MetadataQuery = interpreter.MetadataQuery AccountBalance = interpreter.AccountBalance + ColorBalance = interpreter.ColorBalance Balances = interpreter.Balances AccountMetadata = interpreter.AccountMetadata @@ -76,6 +79,10 @@ type ( MissingFundsErr = interpreter.MissingFundsErr ) +// Uncolored wraps an amount as a ColorBalance under the empty color key +// (the "no color" bucket). Useful when building a Balances literal in tests. +var Uncolored = interpreter.Uncolored + func (p ParseResult) Run(ctx context.Context, vars VariablesMap, store Store) (ExecutionResult, InterpreterError) { return p.RunWithFeatureFlags(ctx, vars, store, nil) } diff --git a/numscript_test.go b/numscript_test.go index b4da92d6..81bdc1b1 100644 --- a/numscript_test.go +++ b/numscript_test.go @@ -128,17 +128,17 @@ send [COIN 100] ( // TODO maybe those calls can be batched together { // this is required by the balance() call - "account_that_needs_balance": {"USD/2"}, + "account_that_needs_balance": {{Asset: "USD/2"}}, }, { // this is defined in the variables - "source1": {"COIN"}, + "source1": {{Asset: "COIN"}}, // this is defined in account metadata - "source2": {"COIN"}, + "source2": {{Asset: "COIN"}}, // this appears as literal - "source3": {"COIN"}, + "source3": {{Asset: "COIN"}}, }, }, store.GetBalancesCalls) @@ -174,8 +174,8 @@ send [COIN 100] ( require.Equal(t, []numscript.BalanceQuery{ { - "a": {"COIN"}, - "b": {"COIN"}, + "a": {{Asset: "COIN"}}, + "b": {{Asset: "COIN"}}, }, }, store.GetBalancesCalls) @@ -205,7 +205,7 @@ func TestDoNotGetBalancesTwice(t *testing.T) { require.Equal(t, []numscript.BalanceQuery{ { - "alice": {"COIN"}, + "alice": {{Asset: "COIN"}}, }, }, store.GetBalancesCalls) @@ -226,8 +226,8 @@ func TestGetBalancesAllotment(t *testing.T) { store := ObservableStore{ StaticStore: interpreter.StaticStore{ Balances: interpreter.Balances{ - "a": {"COIN": big.NewInt(10000)}, - "b": {"COIN": big.NewInt(10000)}, + "a": {"COIN": numscript.Uncolored(big.NewInt(10000))}, + "b": {"COIN": numscript.Uncolored(big.NewInt(10000))}, }, }, } @@ -241,8 +241,8 @@ func TestGetBalancesAllotment(t *testing.T) { require.Equal(t, []numscript.BalanceQuery{ { - "a": {"COIN"}, - "b": {"COIN"}, + "a": {{Asset: "COIN"}}, + "b": {{Asset: "COIN"}}, }, }, store.GetBalancesCalls) @@ -268,7 +268,7 @@ func TestGetBalancesOverdraft(t *testing.T) { require.Equal(t, []numscript.BalanceQuery{ { - "a": {"COIN"}, + "a": {{Asset: "COIN"}}, }, }, store.GetBalancesCalls) @@ -288,7 +288,7 @@ func TestDoNotFetchBalanceTwice(t *testing.T) { require.Equal(t, []numscript.BalanceQuery{ { - "src": {"COIN"}, + "src": {{Asset: "COIN"}}, }, }, store.GetBalancesCalls, @@ -314,10 +314,10 @@ func TestDoNotFetchBalanceTwice2(t *testing.T) { require.Equal(t, []numscript.BalanceQuery{ { - "src1": {"COIN"}, + "src1": {{Asset: "COIN"}}, }, { - "src2": {"COIN"}, + "src2": {{Asset: "COIN"}}, }, }, store.GetBalancesCalls, @@ -343,10 +343,10 @@ func TestDoNotFetchBalanceTwice3(t *testing.T) { require.Equal(t, []numscript.BalanceQuery{ { - "src": {"EUR/2"}, + "src": {{Asset: "EUR/2"}}, }, { - "src": {"USD/2"}, + "src": {{Asset: "USD/2"}}, }, }, store.GetBalancesCalls, @@ -410,7 +410,7 @@ send [USD/2 30] ( require.Equal(t, []numscript.BalanceQuery{ { - "alice": {"USD/2"}, + "alice": {{Asset: "USD/2"}}, }, }, store.GetBalancesCalls, @@ -437,7 +437,7 @@ set_tx_meta( StaticStore: interpreter.StaticStore{ Balances: interpreter.Balances{ "alice": interpreter.AccountBalance{ - "USD/2": big.NewInt(20), + "USD/2": numscript.Uncolored(big.NewInt(20)), }, }, }, @@ -485,10 +485,10 @@ send [USD/2 10] ( }, Balances: interpreter.Balances{ "a": interpreter.AccountBalance{ - "USD/2": big.NewInt(100), + "USD/2": numscript.Uncolored(big.NewInt(100)), }, "a2": interpreter.AccountBalance{ - "USD/2": big.NewInt(1), + "USD/2": numscript.Uncolored(big.NewInt(1)), }, }, }, @@ -519,10 +519,10 @@ send [USD/2 10] ( require.Equal(t, []numscript.BalanceQuery{ { - "a": {"USD/2"}, + "a": {{Asset: "USD/2"}}, }, { - "a2": {"USD/2"}, + "a2": {{Asset: "USD/2"}}, }, }, store.GetBalancesCalls, diff --git a/scripts/migrate-specs-to-color.sh b/scripts/migrate-specs-to-color.sh new file mode 100755 index 00000000..bcc6f40e --- /dev/null +++ b/scripts/migrate-specs-to-color.sh @@ -0,0 +1,61 @@ +#!/usr/bin/env bash +# Migrate .num.specs.json balance entries from the legacy nested shape +# ({"": N, "COLOR": M}) to the value-object shape introduced when `color` +# became a first-class Posting field. +# +# Output rules per (account, asset) entry: +# • uncolored-only → bare number "USD/2": 100 +# • single colored entry → value-object "USD/2": { "color": "RED", "amount": 50 } +# • mixed / multi-color → array of value-objects, sorted by color +# +# Usage: +# scripts/migrate-specs-to-color.sh path/to/file.num.specs.json [more files...] +# scripts/migrate-specs-to-color.sh $(find . -name '*.num.specs.json') +set -euo pipefail + +if [[ $# -eq 0 ]]; then + echo "usage: $0 file.num.specs.json [...]" >&2 + exit 2 +fi + +read -r -d '' JQ_FILTER <<'JQ' || true +def to_value_object: + if type != "object" then . + elif (keys_unsorted | length) == 0 then . + elif (keys_unsorted | length) == 1 and keys_unsorted[0] == "" then + .[""] + elif (keys_unsorted | length) == 1 then + {color: keys_unsorted[0], amount: .[keys_unsorted[0]]} + else + [ (to_entries | sort_by(.key))[] | + if .key == "" then {amount: .value} + else {color: .key, amount: .value} + end ] + end; + +def transform_balances: + if type != "object" then . + else + with_entries(.value |= ( + if type != "object" then . + else with_entries(.value |= to_value_object) + end + )) + end; + +(if has("balances") then .balances |= transform_balances else . end) | +(if has("testCases") then + .testCases |= map( + (if has("balances") then .balances |= transform_balances else . end) | + (if has("expect.endBalances") then ."expect.endBalances" |= transform_balances else . end) | + (if has("expect.endBalances.include") then ."expect.endBalances.include" |= transform_balances else . end) + ) +else . end) +JQ + +for f in "$@"; do + tmp="$(mktemp)" + jq --indent 2 "$JQ_FILTER" "$f" > "$tmp" + # Preserve trailing newline behaviour of jq output. + mv "$tmp" "$f" +done diff --git a/specs.schema.json b/specs.schema.json index a3cd1680..1a55424b 100644 --- a/specs.schema.json +++ b/specs.schema.json @@ -92,13 +92,31 @@ "additionalProperties": false, "patternProperties": { "^([A-Z]+(/[0-9]+)?)$": { - "type": "number" + "oneOf": [ + { "type": "number", "description": "Uncolored shorthand" }, + { "$ref": "#/definitions/BalanceEntry" }, + { "type": "array", "items": { "$ref": "#/definitions/BalanceEntry" } } + ] } } } } }, + "BalanceEntry": { + "type": "object", + "additionalProperties": false, + "required": ["amount"], + "properties": { + "color": { + "type": "string", + "pattern": "^[A-Z]*$", + "description": "Color bucket. Omit (or set to \"\") for the uncolored bucket." + }, + "amount": { "type": "number" } + } + }, + "VariablesMap": { "type": "object", "description": "Map of variable name to variable stringified value",