Color of money base implementation

Lexer.g4

@@ -31,6 +31,7 @@ STAR: '*';
 PLUS: '+';
 MINUS: '-';
 DIV: '/';
+RESTRICT: '\\';
 
 PERCENTAGE_PORTION_LITERAL: [0-9]+ ('.' [0-9]+)? '%';
 

Numscript.g4

@@ -41,11 +41,13 @@ allotment:
 	valueExpr	# portionedAllotment
 	| REMAINING	# remainingAllotment;
 
+colorConstraint: RESTRICT valueExpr;
+
 source:
-	address = valueExpr ALLOWING UNBOUNDED OVERDRAFT						# srcAccountUnboundedOverdraft
-	| address = valueExpr ALLOWING OVERDRAFT UP TO maxOvedraft = valueExpr	#
+	address = valueExpr colorConstraint? ALLOWING UNBOUNDED OVERDRAFT						# srcAccountUnboundedOverdraft
+	| address = valueExpr colorConstraint? ALLOWING OVERDRAFT UP TO maxOvedraft = valueExpr	#
 		srcAccountBoundedOverdraft
-	| valueExpr							# srcAccount
+	| valueExpr colorConstraint?		# srcAccount
 	| LBRACE allotmentClauseSrc+ RBRACE	# srcAllotment
 	| LBRACE source* RBRACE				# srcInorder
 	| ONEOF LBRACE source+ RBRACE		# srcOneof

internal/analysis/check.go

@@ -476,18 +476,29 @@ func (res *CheckResult) checkSource(source parser.Source) {
 	switch source := source.(type) {
 	case *parser.SourceAccount:
 		res.checkExpression(source.ValueExpr, TypeAccount)
+		res.checkExpression(source.Color, TypeString)
 		if account, ok := source.ValueExpr.(*parser.AccountInterpLiteral); ok {
 			if account.IsWorld() && res.unboundedSend {
 				res.pushDiagnostic(source.GetRange(), InvalidUnboundedAccount{})
 			} else if account.IsWorld() {
 				res.unboundedAccountInSend = account
 			}
 
-			if _, emptied := res.emptiedAccount[account.String()]; emptied && !account.IsWorld() {
+			coloredAccountName := account.String()
+			switch col := source.Color.(type) {
+			case *parser.Variable:
+				coloredAccountName += "\\$" + col.Name
+			case *parser.StringLiteral:
+				if col.String != "" {
+					coloredAccountName += "\\\"" + col.String + "\""
+				}
+			}
+
+			if _, emptied := res.emptiedAccount[coloredAccountName]; emptied && !account.IsWorld() {
 				res.pushDiagnostic(account.Range, EmptiedAccount{Name: account.String()})
 			}
 
-			res.emptiedAccount[account.String()] = struct{}{}
+			res.emptiedAccount[coloredAccountName] = struct{}{}
 		}
 
 	case *parser.SourceOverdraft:
@@ -507,6 +518,7 @@ func (res *CheckResult) checkSource(source parser.Source) {
 		}
 
 		res.checkExpression(source.Address, TypeAccount)
+		res.checkExpression(source.Color, TypeString)
 		if source.Bounded != nil {
 			res.checkExpression(*source.Bounded, TypeMonetary)
 		}

internal/analysis/check_test.go

@@ -861,3 +861,69 @@ send $m ( source = @world destination = @dest)
 		checkSource(input),
 	)
 }
+
+func TestCheckColorVar(t *testing.T) {
+	t.Parallel()
+
+	input := `
+	send [COIN 100] (
+  	source = {
+			@a \ $col1
+			@b \ $col2 allowing unbounded overdraft
+		}
+  	destination = @dest
+)`
+
+	require.Equal(t,
+		[]analysis.Diagnostic{
+			{
+				Range: parser.RangeOfIndexed(input, "$col1", 0),
+				Kind:  analysis.UnboundVariable{Name: "col1", Type: "string"},
+			},
+			{
+				Range: parser.RangeOfIndexed(input, "$col2", 0),
+				Kind:  analysis.UnboundVariable{Name: "col2", Type: "string"},
+			},
+		},
+		checkSource(input),
+	)
+}
+
+func TestInorderRedundantWhenColored(t *testing.T) {
+	t.Parallel()
+
+	input := `send [COIN 100] (
+  	source = {
+			@a
+			@a \ "x"
+		}
+  	destination = @dest
+)`
+
+	require.Equal(t,
+		[]analysis.Diagnostic(nil),
+		checkSource(input),
+	)
+}
+
+func TestInorderRedundantWhenEmptyColored(t *testing.T) {
+	t.Parallel()
+
+	input := `send [COIN 100] (
+  	source = {
+			@a
+			@a \ "" // <- empty color behaves as no color
+		}
+  	destination = @dest
+)`
+
+	require.Equal(t,
+		[]analysis.Diagnostic{
+			{
+				Range: parser.RangeOfIndexed(input, "@a", 1),
+				Kind:  analysis.EmptiedAccount{Name: "a"},
+			},
+		},
+		checkSource(input),
+	)
+}

internal/flags/flags.go

@@ -9,6 +9,7 @@ const (
 	ExperimentalOneofFeatureFlag             FeatureFlag = "experimental-oneof"
 	ExperimentalAccountInterpolationFlag     FeatureFlag = "experimental-account-interpolation"
 	ExperimentalMidScriptFunctionCall        FeatureFlag = "experimental-mid-script-function-call"
+	ExperimentalAssetColors                  FeatureFlag = "experimental-asset-colors"
 )
 
 var AllFlags []string = []string{
@@ -18,4 +19,5 @@ var AllFlags []string = []string{
 	ExperimentalOneofFeatureFlag,
 	ExperimentalAccountInterpolationFlag,
 	ExperimentalMidScriptFunctionCall,
+	ExperimentalAssetColors,
 }

internal/interpreter/balances.go

@@ -2,6 +2,7 @@ package interpreter
 
 import (
 	"math/big"
+	"strings"
 
 	"github.com/formancehq/numscript/internal/utils"
 )
@@ -12,6 +13,21 @@ func (b Balances) fetchAccountBalances(account string) AccountBalance {
 	})
 }
 
+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, asset string) *big.Int {

internal/interpreter/batch_balances_query.go

@@ -31,7 +31,7 @@
 		if err != nil {
 			return err
 		}
-		st.batchQuery(*account, *asset)
+		st.batchQuery(*account, *asset, nil)
 		return nil
 
 	case *parser.SendStatement:
@@ -50,10 +50,11 @@
 	}
 }
 
-func (st *programState) batchQuery(account string, asset string) {
+func (st *programState) batchQuery(account string, asset string, color *string) {
 	if account == "world" {
 		return
 	}
+	asset = coloredAsset(asset, color)
 
 	previousValues := st.CurrentBalanceQuery[account]
 	if !slices.Contains(previousValues, asset) {
@@ -89,7 +90,12 @@
 			return err
 		}
 
-		st.batchQuery(*account, st.CurrentAsset)
+		color, err := evaluateOptExprAs(st, source.Color, expectString)
+		if err != nil {
+			return err
+		}
+
+		st.batchQuery(*account, st.CurrentAsset, color)
 		return nil
 
 	case *parser.SourceOverdraft:
@@ -102,7 +108,12 @@
 		if err != nil {
 			return err
 		}
-		st.batchQuery(*account, st.CurrentAsset)
+		color, err := evaluateOptExprAs(st, source.Color, expectString)
+		if err != nil {
+			return err
+		}
+
+		st.batchQuery(*account, st.CurrentAsset, color)
 		return nil
 
 	case *parser.SourceInorder:

internal/interpreter/evaluate_expr.go

@@ -101,6 +101,13 @@
 	}
 }
 
+func evaluateOptExprAs[T any](st *programState, expr parser.ValueExpr, expect func(Value, parser.Range) (*T, InterpreterError)) (*T, InterpreterError) {
+	if expr == nil {
+		return nil, nil
+	}
+	return evaluateExprAs(st, expr, expect)
+}
+
 func evaluateExprAs[T any](st *programState, expr parser.ValueExpr, expect func(Value, parser.Range) (*T, InterpreterError)) (*T, InterpreterError) {
 	value, err := st.evaluateExpr(expr)
 	if err != nil {
@@ -127,6 +134,26 @@
 	return values, nil
 }
 
+func (s *programState) evaluateColor(colorExpr parser.ValueExpr) (*string, InterpreterError) {
+	color, err := evaluateOptExprAs(s, colorExpr, expectString)
+	if err != nil {
+		return nil, err
+	}
+	if color == nil {
+		return nil, nil
+	}
+
+	isValidColor := colorRe.Match([]byte(*color))
+	if !isValidColor {
+		return nil, InvalidColor{
+			Range: colorExpr.GetRange(),
+			Color: *color,
+		}
+	}
+
+	return color, nil
+}
+
 func (st *programState) plusOp(left parser.ValueExpr, right parser.ValueExpr) (Value, InterpreterError) {
 	leftValue, err := evaluateExprAs(st, left, expectOneOf(
 		expectMapped(expectMonetary, func(m Monetary) opAdd {