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

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

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

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