Feat/expr in var origin (#45)

Author: ascandone

Numscript.g4

@@ -21,13 +21,14 @@ valueExpr:
 	| monetaryLit													# monetaryLiteral
 	| left = valueExpr op = DIV right = valueExpr					# infixExpr
 	| left = valueExpr op = (PLUS | MINUS) right = valueExpr		# infixExpr
-	| LPARENS valueExpr RPARENS										# parenthesizedExpr;
+	| LPARENS valueExpr RPARENS										# parenthesizedExpr
+	| functionCall													# application;
 
 functionCallArgs: valueExpr ( COMMA valueExpr)*;
 functionCall:
 	fnName = (OVERDRAFT | IDENTIFIER) LPARENS functionCallArgs? RPARENS;
 
-varOrigin: EQ functionCall;
+varOrigin: EQ valueExpr;
 varDeclaration:
 	type_ = IDENTIFIER name = VARIABLE_NAME varOrigin?;
 varsDeclaration: VARS LBRACE varDeclaration* RBRACE;

go.mod

@@ -32,5 +32,5 @@ require (
 	github.com/tidwall/match v1.1.1 // indirect
 	github.com/tidwall/pretty v1.2.1 // indirect
 	github.com/tidwall/sjson v1.2.5 // indirect
-	golang.org/x/exp v0.0.0-20240707233637-46b078467d37
+	golang.org/x/exp v0.0.0-20240707233637-46b078467d37 // indirect
 )

internal/analysis/check.go

@@ -162,7 +162,7 @@ func (res *CheckResult) check() {
 			}
 
 			if varDecl.Origin != nil {
-				res.checkVarOrigin(*varDecl.Origin, varDecl)
+				res.checkExpression(*varDecl.Origin, varDecl.Type.Name)
 			}
 		}
 	}
@@ -302,18 +302,20 @@ func (res *CheckResult) checkDuplicateVars(variableName parser.Variable, decl pa
 	}
 }
 
-func (res *CheckResult) checkVarOrigin(fnCall parser.FnCall, decl parser.VarDeclaration) {
-	resolution, ok := Builtins[fnCall.Caller.Name]
-	if ok {
-		resolution, ok := resolution.(VarOriginFnCallResolution)
-		if ok {
-			res.fnCallResolution[decl.Origin.Caller] = resolution
-			res.assertHasType(decl.Name, resolution.Return, decl.Type.Name)
+func (res *CheckResult) checkFnCall(fnCall parser.FnCall) string {
+	returnType := TypeAny
+
+	if resolution, ok := Builtins[fnCall.Caller.Name]; ok {
+		if resolution, ok := resolution.(VarOriginFnCallResolution); ok {
+			res.fnCallResolution[fnCall.Caller] = resolution
+			returnType = resolution.Return
 		}
 	}
 
 	// this must come after resolution
 	res.checkFnCallArity(&fnCall)
+
+	return returnType
 }
 
 func (res *CheckResult) checkExpression(lit parser.ValueExpr, requiredType string) {
@@ -379,6 +381,9 @@ func (res *CheckResult) checkTypeOf(lit parser.ValueExpr, typeHint string) strin
 	case *parser.StringLiteral:
 		return TypeString
 
+	case *parser.FnCall:
+		return res.checkFnCall(*lit)
+
 	default:
 		return TypeAny
 	}

internal/analysis/hover.go

@@ -75,7 +75,7 @@ func hoverOnVar(varDecl parser.VarDeclaration, position parser.Position) Hover {
 	}
 
 	if varDecl.Origin != nil {
-		hover := hoverOnFnCall(*varDecl.Origin, position)
+		hover := hoverOnExpression(*varDecl.Origin, position)
 		if hover != nil {
 			return hover
 		}
@@ -183,6 +183,8 @@ func hoverOnExpression(lit parser.ValueExpr, position parser.Position) Hover {
 			return hover
 		}
 
+	case *parser.FnCall:
+		return hoverOnFnCall(*lit, position)
 	}
 
 	return nil

internal/cmd/run.go

@@ -29,6 +29,7 @@ var runOutFormatOpt string
 var overdraftFeatureFlag bool
 var oneOfFeatureFlag bool
 var accountInterpolationFlag bool
+var midScriptFunctionCallFeatureFlag bool
 
 type inputOpts struct {
 	Script    string                       `json:"script"`
@@ -129,6 +130,9 @@ func run(path string) {
 	if accountInterpolationFlag {
 		featureFlags[interpreter.ExperimentalAccountInterpolationFlag] = struct{}{}
 	}
+	if midScriptFunctionCallFeatureFlag {
+		featureFlags[interpreter.ExperimentalMidScriptFunctionCall] = struct{}{}
+	}
 
 	result, err := interpreter.RunProgram(context.Background(), parseResult.Value, opt.Variables, interpreter.StaticStore{
 		Balances: opt.Balances,
@@ -208,9 +212,11 @@ func getRunCmd() *cobra.Command {
 	cmd.Flags().BoolVar(&runStdinFlag, "stdin", false, "Take input from stdin (same format as the --raw option)")
 
 	// Feature flag
-	cmd.Flags().BoolVar(&overdraftFeatureFlag, interpreter.ExperimentalOverdraftFunctionFeatureFlag, false, "feature flag to enable the overdraft() function")
-	cmd.Flags().BoolVar(&oneOfFeatureFlag, interpreter.ExperimentalOneofFeatureFlag, false, "feature flag to enable the oneof combinator")
+
+	cmd.Flags().BoolVar(&overdraftFeatureFlag, interpreter.ExperimentalOverdraftFunctionFeatureFlag, false, "enables the experimental overdraft() function")
+	cmd.Flags().BoolVar(&oneOfFeatureFlag, interpreter.ExperimentalOneofFeatureFlag, false, "enable the experimental oneof combinator")
 	cmd.Flags().BoolVar(&accountInterpolationFlag, interpreter.ExperimentalAccountInterpolationFlag, false, "enables an account interpolation syntax, e.g. @users:$id:pending")
+	cmd.Flags().BoolVar(&midScriptFunctionCallFeatureFlag, interpreter.ExperimentalMidScriptFunctionCall, false, "allows to use function call as expression, and to use any expression when definining variables")
 
 	// Output options
 	cmd.Flags().StringVar(&runOutFormatOpt, "output-format", OutputFormatPretty, "Set the output format. Available options: pretty, json.")

internal/interpreter/balances.go

@@ -0,0 +1,63 @@
+package interpreter
+
+import (
+	"math/big"
+
+	"github.com/formancehq/numscript/internal/utils"
+)
+
+func (b Balances) fetchAccountBalances(account string) AccountBalance {
+	return defaultMapGet(b, account, func() AccountBalance {
+		return AccountBalance{}
+	})
+}
+
+// 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 {
+	accountBalances := b.fetchAccountBalances(account)
+
+	return defaultMapGet(accountBalances, asset, func() *big.Int {
+		return new(big.Int)
+	})
+}
+
+func (b Balances) has(account string, asset string) bool {
+	accountBalances := defaultMapGet(b, account, func() AccountBalance {
+		return AccountBalance{}
+	})
+
+	_, ok := accountBalances[asset]
+	return ok
+}
+
+// given a BalanceQuery, return a new query which only contains needed (asset, account) 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)
+		})
+
+		if len(filteredCurrencies) > 0 {
+			filteredQuery[accountName] = filteredCurrencies
+		}
+
+	}
+	return filteredQuery
+}
+
+// Merge balances by adding balances in the "update" arg
+func (b Balances) mergeBalance(update Balances) {
+	// merge queried balance
+	for acc, accBalances := range update {
+		cachedAcc := defaultMapGet(b, acc, func() AccountBalance {
+			return AccountBalance{}
+		})
+
+		for curr, amt := range accBalances {
+			cachedAcc[curr] = amt
+		}
+	}
+}

internal/interpreter/balances_test.go

@@ -0,0 +1,31 @@
+package interpreter
+
+import (
+	"math/big"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestFilterQuery(t *testing.T) {
+	fullBalance := Balances{
+		"alice": AccountBalance{
+			"EUR/2": big.NewInt(1),
+			"USD/2": big.NewInt(2),
+		},
+		"bob": AccountBalance{
+			"BTC": big.NewInt(3),
+		},
+	}
+
+	filteredQuery := fullBalance.filterQuery(BalanceQuery{
+		"alice":   []string{"GBP/2", "YEN", "EUR/2"},
+		"bob":     []string{"BTC"},
+		"charlie": []string{"ETH"},
+	})
+
+	require.Equal(t, BalanceQuery{
+		"alice":   []string{"GBP/2", "YEN"},
+		"charlie": []string{"ETH"},
+	}, filteredQuery)
+}

internal/interpreter/batch_balances_query.go

@@ -5,7 +5,6 @@ import (
 
 	"github.com/formancehq/numscript/internal/parser"
 	"github.com/formancehq/numscript/internal/utils"
-	"golang.org/x/exp/maps"
 )
 
 // traverse the script to batch in advance required balance queries
@@ -57,41 +56,28 @@ func (st *programState) batchQuery(account string, asset string) {
 	}
 
 	previousValues := st.CurrentBalanceQuery[account]
-	if !slices.Contains[[]string, string](previousValues, asset) {
+	if !slices.Contains(previousValues, asset) {
 		st.CurrentBalanceQuery[account] = append(previousValues, asset)
 	}
 }
 
 func (st *programState) runBalancesQuery() error {
-	filteredQuery := BalanceQuery{}
-	for accountName, queriedCurrencies := range st.CurrentBalanceQuery {
-
-		cachedCurrenciesForAccount := defaultMapGet(st.CachedBalances, accountName, func() AccountBalance {
-			return AccountBalance{}
-		})
-
-		for _, queriedCurrency := range queriedCurrencies {
-			isAlreadyCached := slices.Contains(maps.Keys(cachedCurrenciesForAccount), queriedCurrency)
-			if !isAlreadyCached {
-				filteredQuery[accountName] = queriedCurrencies
-			}
-		}
-
-	}
+	filteredQuery := st.CachedBalances.filterQuery(st.CurrentBalanceQuery)
 
 	// avoid updating balances if we don't need to fetch new data
 	if len(filteredQuery) == 0 {
 		return nil
 	}
 
-	balances, err := st.Store.GetBalances(st.ctx, filteredQuery)
+	queriedBalances, err := st.Store.GetBalances(st.ctx, filteredQuery)
 	if err != nil {
 		return err
 	}
 	// reset batch query
 	st.CurrentBalanceQuery = BalanceQuery{}
 
-	st.CachedBalances = balances
+	st.CachedBalances.mergeBalance(queriedBalances)
+
 	return nil
 }
 

internal/interpreter/evaluate_expr.go

@@ -84,6 +84,16 @@ func (st *programState) evaluateExpr(expr parser.ValueExpr) (Value, InterpreterE
 			return nil, nil
 		}
 
+	case *parser.FnCall:
+		if !st.varOriginPosition {
+			err := st.checkFeatureFlag(ExperimentalMidScriptFunctionCall)
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		return st.handleFnCall(nil, *expr)
+
 	default:
 		utils.NonExhaustiveMatchPanic[any](expr)
 		return nil, nil