feat!: expose color as a first-class Posting property

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)
 		}
 	}
 

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)
 }

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       |
 ---

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
 }

internal/interpreter/balances.go

@@ -1,111 +1,186 @@
 package interpreter
 
 import (
+	"encoding/json"
+	"fmt"
 	"math/big"
-	"strings"
 
 	"github.com/formancehq/numscript/internal/utils"
 )
 
-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)
-			})
+// 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}
+}
+
+// UnmarshalJSON accepts two JSON shapes for a Balances value:
+//
+//  1. Flat (uncolored shorthand):
+//     {"alice": {"USD/2": 100, "EUR/2": -42}}
+//     Every asset gets a single entry under the "" (no color) bucket.
+//
+//  2. Colored (full form):
+//     {"alice": {"USD/2": {"": 100, "GRANTS": 50}}}
+//
+// Shapes can be mixed across assets within the same document.
+func (b *Balances) UnmarshalJSON(data []byte) error {
+	raw := map[string]map[string]json.RawMessage{}
+	if err := json.Unmarshal(data, &raw); err != nil {
+		return err
+	}
+
+	out := Balances{}
+	for account, assets := range raw {
+		accB := AccountBalance{}
+		out[account] = accB
+		for asset, rawVal := range assets {
+			colorMap, err := decodeColorBalance(rawVal)
+			if err != nil {
+				return fmt.Errorf("balances[%s][%s]: %w", account, asset, err)
+			}
+			accB[asset] = colorMap
 		}
 	}
-	return cloned
+	*b = out
+	return nil
 }
 
-func coloredAsset(asset string, color *string) string {
-	if color == nil || *color == "" {
-		return asset
+func decodeColorBalance(data json.RawMessage) (ColorBalance, error) {
+	// Try the shorthand: a single number meaning the uncolored bucket.
+	var amount json.Number
+	if err := json.Unmarshal(data, &amount); err == nil {
+		n, ok := new(big.Int).SetString(amount.String(), 10)
+		if !ok {
+			return nil, fmt.Errorf("invalid integer amount %q", amount.String())
+		}
+		return ColorBalance{"": n}, nil
 	}
 
-	// note: 1 <= len(parts) <= 2
-	parts := strings.Split(asset, "/")
-
-	coloredAsset := parts[0] + "_" + *color
-	if len(parts) > 1 {
-		coloredAsset += "/" + parts[1]
+	// Otherwise expect a {color: amount} object.
+	raw := map[string]json.Number{}
+	if err := json.Unmarshal(data, &raw); err != nil {
+		return nil, fmt.Errorf("expected integer or {color: amount} object, got %s", string(data))
+	}
+	out := ColorBalance{}
+	for color, amt := range raw {
+		n, ok := new(big.Int).SetString(amt.String(), 10)
+		if !ok {
+			return nil, fmt.Errorf("color %q: invalid integer amount %q", color, amt.String())
+		}
+		out[color] = n
 	}
-	return coloredAsset
+	return out, nil
 }
 
-// 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 {
-		return new(big.Int)
-	})
+func (b Balances) DeepClone() Balances {
+	cloned := make(Balances)
+	for account, accountBalances := range b {
+		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 (b Balances) has(account string, asset string) bool {
-	accountBalances := utils.MapGetOrPutDefault(b, account, func() AccountBalance {
+// 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)
+	})
+}
 
-	_, 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
+			})
 		})
 	})
 }