feat: asset scaling

Lexer.g4

@@ -32,6 +32,8 @@ PLUS: '+';
 MINUS: '-';
 DIV: '/';
 RESTRICT: '\\';
+WITH: 'with';
+SCALING: 'scaling';
 
 PERCENTAGE_PORTION_LITERAL: [0-9]+ ('.' [0-9]+)? '%';
 

Numscript.g4

@@ -43,10 +43,10 @@ allotment:
 
 colorConstraint: RESTRICT valueExpr;
 
-source:
-	address = valueExpr colorConstraint? ALLOWING UNBOUNDED OVERDRAFT						# srcAccountUnboundedOverdraft
-	| address = valueExpr colorConstraint? ALLOWING OVERDRAFT UP TO maxOvedraft = valueExpr	#
-		srcAccountBoundedOverdraft
+source
+	: address = valueExpr colorConstraint? ALLOWING UNBOUNDED OVERDRAFT						# srcAccountUnboundedOverdraft
+	| address = valueExpr colorConstraint? ALLOWING OVERDRAFT UP TO maxOvedraft = valueExpr	#srcAccountBoundedOverdraft
+	| address = valueExpr colorConstraint? WITH SCALING						# srcAccountWithScaling
 	| valueExpr colorConstraint?		# srcAccount
 	| LBRACE allotmentClauseSrc+ RBRACE	# srcAllotment
 	| LBRACE source* RBRACE				# srcInorder

internal/analysis/check.go

@@ -593,6 +593,10 @@ func (res *CheckResult) checkSource(source parser.Source) {
 			res.unifyNodeWith(*source.Bounded, res.stmtType)
 		}
 
+	case *parser.SourceWithScaling:
+		res.checkExpression(source.Address, TypeAccount)
+		res.checkExpression(source.Color, TypeString)
+
 	case *parser.SourceInorder:
 		for _, source := range source.Sources {
 			res.checkSource(source)

internal/interpreter/asset_scaling.go

@@ -0,0 +1,133 @@
+package interpreter
+
+import (
+	"fmt"
+	"math/big"
+	"slices"
+	"strconv"
+	"strings"
+
+	"github.com/formancehq/numscript/internal/utils"
+)
+
+func assetToScaledAsset(asset string) string {
+	// GPT-generated TODO double check
+	parts := strings.Split(asset, "/")
+	if len(parts) == 1 {
+		return asset + "/*"
+	}
+	return parts[0] + "/*"
+}
+
+func buildScaledAsset(baseAsset string, scale int64) string {
+	if scale == 0 {
+		return baseAsset
+	}
+	return fmt.Sprintf("%s/%d", baseAsset, scale)
+}
+
+func getAssetScale(asset string) (string, int64) {
+	parts := strings.Split(asset, "/")
+	if len(parts) == 2 {
+		scale, err := strconv.ParseInt(parts[1], 10, 64)
+		if err == nil {
+			return parts[0], scale
+		}
+		// fallback if parsing fails
+		return parts[0], 0
+	}
+	return asset, 0
+}
+
+func getAssets(balance AccountBalance, baseAsset string) map[int64]*big.Int {
+	result := make(map[int64]*big.Int)
+	for asset, amount := range balance {
+		if strings.HasPrefix(asset, baseAsset) {
+			_, scale := getAssetScale(asset)
+			result[scale] = amount
+		}
+	}
+	return result
+}
+
+// e.g.
+//
+// need=[EUR/2 100], got={EUR/2: 100, EUR: 1}
+// => {EUR/2: 100, EUR: 1}
+//
+// need=[EUR 1], got={EUR/2: 100, EUR: 0}
+// => {EUR/2: 100, EUR: 0}
+//
+// need=[EUR/2 199], got={EUR/2: 100, EUR: 2}
+// => {EUR/2: 100, EUR: 1}
+func findSolution(
+	neededAmt *big.Int,
+	neededAmtScale int64,
+	scales map[int64]*big.Int,
+) map[int64]*big.Int {
+	// we clone neededAmt so that we can update it
+	neededAmt = new(big.Int).Set(neededAmt)
+
+	type scalePair struct {
+		scale  int64
+		amount *big.Int
+	}
+
+	var assets []scalePair
+	for k, v := range scales {
+		assets = append(assets, scalePair{
+			scale:  k,
+			amount: v,
+		})
+	}
+
+	// Sort in ASC order (e.g. EUR, EUR/2, ..)
+	slices.SortFunc(assets, func(p scalePair, other scalePair) int {
+		return int(p.scale - other.scale)
+	})
+
+	out := map[int64]*big.Int{}
+
+	left := new(big.Int).Set(neededAmt)
+
+	for _, p := range assets {
+		scaleDiff := neededAmtScale - p.scale
+
+		exp := big.NewInt(scaleDiff)
+		exp.Abs(exp)
+		exp.Exp(big.NewInt(10), exp, nil)
+
+		// scalingFactor := 10 ^ (neededAmtScale - p.scale)
+		// note that 10^0 == 1 and 10^(-n) == 1/(10^n)
+		scalingFactor := new(big.Rat).SetInt(exp)
+		if scaleDiff < 0 {
+			scalingFactor.Inv(scalingFactor)
+		}
+
+		allowed := new(big.Int).Mul(p.amount, scalingFactor.Num())
+		allowed.Div(allowed, scalingFactor.Denom())
+
+		taken := utils.MinBigInt(allowed, neededAmt)
+
+		intPart := new(big.Int).Mul(taken, scalingFactor.Denom())
+		intPart.Div(intPart, scalingFactor.Num())
+
+		if intPart.Cmp(big.NewInt(0)) == 0 {
+			continue
+		}
+
+		neededAmt.Sub(neededAmt, taken)
+		out[p.scale] = intPart
+
+		// if neededAmt <= 0
+		if neededAmt.Cmp(big.NewInt(0)) != 1 {
+			return out
+		}
+	}
+
+	if left.Cmp(big.NewInt(0)) != 0 {
+		return nil
+	}
+
+	return out
+}

internal/interpreter/asset_scaling_test.go

@@ -0,0 +1,82 @@
+package interpreter
+
+import (
+	"math/big"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestScalingZeroNeeded(t *testing.T) {
+	t.Skip()
+
+	// need [EUR/2 ]
+	sol := findSolution(
+		big.NewInt(0),
+		42,
+		map[int64]*big.Int{
+			2: big.NewInt(100),
+			1: big.NewInt(1),
+		})
+
+	require.Equal(t, map[int64]*big.Int{
+		42: big.NewInt(0),
+	}, sol)
+}
+
+func TestScalingSameAsset(t *testing.T) {
+	sol := findSolution(
+		// Need [EUR/2 200]
+		big.NewInt(200),
+		2,
+
+		// Have: {EUR/2: 201}
+		map[int64]*big.Int{
+			2: big.NewInt(201),
+		})
+
+	require.Equal(t, map[int64]*big.Int{
+		2: big.NewInt(200),
+	}, sol)
+}
+
+func TestScalingSolutionLowerScale(t *testing.T) {
+	sol := findSolution(
+		big.NewInt(1),
+		0,
+		map[int64]*big.Int{
+			2: big.NewInt(900),
+		})
+
+	require.Equal(t, map[int64]*big.Int{
+		2: big.NewInt(100),
+	}, sol)
+}
+
+func TestScalingSolutionHigherScale(t *testing.T) {
+	sol := findSolution(
+		// Need [EUR/2 200]
+		big.NewInt(200),
+		2,
+
+		// Have: {EUR: 4} (eq to EUR/2 400)
+		map[int64]*big.Int{
+			0: big.NewInt(4),
+		})
+
+	require.Equal(t, map[int64]*big.Int{
+		0: big.NewInt(2),
+	}, sol)
+}
+
+func TestScalingSolutionHigherScaleNoSolution(t *testing.T) {
+	sol := findSolution(
+		big.NewInt(1),
+		2,
+		map[int64]*big.Int{
+			0: big.NewInt(100),
+			1: big.NewInt(100),
+		})
+
+	require.Nil(t, sol)
+}

internal/interpreter/batch_balances_query.go

@@ -98,6 +98,20 @@ func (st *programState) findBalancesQueries(source parser.Source) InterpreterErr
 		st.batchQuery(*account, st.CurrentAsset, color)
 		return nil
 
+	case *parser.SourceWithScaling:
+		account, err := evaluateExprAs(st, source.Address, expectAccount)
+		if err != nil {
+			return err
+		}
+
+		color, err := evaluateOptExprAs(st, source.Color, expectString)
+		if err != nil {
+			return err
+		}
+
+		st.batchQuery(*account, assetToScaledAsset(st.CurrentAsset), color)
+		return nil
+
 	case *parser.SourceOverdraft:
 		// Skip balance tracking when balance is overdraft
 		if source.Bounded == nil {

internal/interpreter/interpreter.go

@@ -2,6 +2,7 @@ package interpreter
 
 import (
 	"context"
+	"fmt"
 	"math/big"
 	"regexp"
 	"strings"
@@ -51,12 +52,26 @@ func (s StaticStore) GetBalances(_ context.Context, q BalanceQuery) (Balances, e
 		})
 
 		for _, curr := range queriedCurrencies {
-			n := new(big.Int)
-			outputAccountBalance[curr] = n
-
-			if i, ok := accountBalanceLookup[curr]; ok {
-				n.Set(i)
+			baseAsset, isCatchAll := strings.CutSuffix(curr, "/*")
+			if isCatchAll {
+
+				for k, v := range accountBalanceLookup {
+					matchesAsset := k == baseAsset || strings.HasPrefix(k, baseAsset+"/")
+					if !matchesAsset {
+						continue
+					}
+					outputAccountBalance[k] = new(big.Int).Set(v)
+				}
+
+			} else {
+				n := new(big.Int)
+				outputAccountBalance[curr] = n
+
+				if i, ok := accountBalanceLookup[curr]; ok {
+					n.Set(i)
+				}
 			}
+
 		}
 	}
 
@@ -507,6 +522,9 @@ func (s *programState) sendAll(source parser.Source) (*big.Int, InterpreterError
 		}
 		return s.sendAllToAccount(source.Address, cap, source.Color)
 
+	case *parser.SourceWithScaling:
+		panic("TODO implement")
+
 	case *parser.SourceInorder:
 		totalSent := big.NewInt(0)
 		for _, subSource := range source.Sources {
@@ -617,6 +635,57 @@ func (s *programState) trySendingUpTo(source parser.Source, amount *big.Int) (*b
 	case *parser.SourceAccount:
 		return s.trySendingToAccount(source.ValueExpr, amount, big.NewInt(0), source.Color)
 
+	case *parser.SourceWithScaling:
+		account, err := evaluateExprAs(s, source.Address, expectAccount)
+		if err != nil {
+			return nil, err
+		}
+
+		baseAsset, assetScale := getAssetScale(s.CurrentAsset)
+		acc, ok := s.CachedBalances[*account]
+		if !ok {
+			panic("TODO accountBal not found")
+		}
+
+		sol := findSolution(
+			amount,
+			assetScale,
+			getAssets(acc, baseAsset),
+		)
+
+		if sol == nil {
+			// we already know we are failing, but we're delegating to the "standard" (non-scaled) mode
+			// so that we get a somewhat helpful (although limited) error message
+			return s.trySendingToAccount(source.Address, amount, big.NewInt(0), source.Color)
+		}
+
+		for scale, sending := range sol {
+			// here we manually emit postings based on the known solution,
+			// and update balances accordingly
+			asset := buildScaledAsset(baseAsset, scale)
+			s.Postings = append(s.Postings, Posting{
+				Source:      *account,
+				Destination: fmt.Sprintf("%s:scaling", *account),
+				Amount:      new(big.Int).Set(sending),
+				Asset:       asset,
+			})
+			acc[asset].Sub(acc[asset], sending)
+		}
+
+		s.Postings = append(s.Postings, Posting{
+			Source:      fmt.Sprintf("%s:scaling", *account),
+			Destination: *account,
+			Amount:      new(big.Int).Set(amount),
+			Asset:       s.CurrentAsset,
+		})
+
+		accBalance := utils.MapGetOrPutDefault(acc, s.CurrentAsset, func() *big.Int {
+			return big.NewInt(0)
+		})
+		accBalance.Add(accBalance, amount)
+
+		return s.trySendingToAccount(source.Address, amount, big.NewInt(0), source.Color)
+
 	case *parser.SourceOverdraft:
 		var cap *big.Int
 		if source.Bounded != nil {