implemented save statement

internal/interpreter/batch_balances_query.go

@@ -15,26 +15,32 @@ func (st *programState) findBalancesQueriesInStatement(statement parser.Statemen
 	case *parser.FnCall:
 		return nil
 
-	case *parser.SendStatement:
-		// set the current asset
-		switch sentValue := statement.SentValue.(type) {
-		case *parser.SentValueAll:
-			asset, err := evaluateLitExpecting(st, sentValue.Asset, expectAsset)
-			if err != nil {
-				return err
-			}
-			st.CurrentAsset = *asset
+	case *parser.SaveStatement:
+		asset, _, err := st.evaluateSentAmt(statement.SentValue)
+		if err != nil {
+			return err
+		}
 
-		case *parser.SentValueLiteral:
-			monetary, err := evaluateLitExpecting(st, sentValue.Monetary, expectMonetary)
-			if err != nil {
-				return err
-			}
-			st.CurrentAsset = string(monetary.Asset)
+		// Although we don't technically need this account's balance rn,
+		// having access to the balance simplifies the "save" statement implementation
+		// this means that we would have a needless query in the case in which the account
+		// which is selected in the "save" statement never actually appears as source
+		//
+		// this would mean that the "save" statement was not needed in the first place,
+		// so preventing this query would hardly be an useful optimization
+		account, err := evaluateLitExpecting(st, statement.Literal, expectAccount)
+		if err != nil {
+			return err
+		}
+		st.batchQuery(*account, *asset)
+		return nil
 
-		default:
-			utils.NonExhaustiveMatchPanic[any](sentValue)
+	case *parser.SendStatement:
+		asset, _, err := st.evaluateSentAmt(statement.SentValue)
+		if err != nil {
+			return err
 		}
+		st.CurrentAsset = *asset
 
 		// traverse source
 		return st.findBalancesQueries(statement.Source)
@@ -51,7 +57,7 @@ func (st *programState) batchQuery(account string, asset string) {
 	}
 
 	previousValues := st.CurrentBalanceQuery[account]
-	if !slices.Contains[[]string, string](previousValues, account) {
+	if !slices.Contains[[]string, string](previousValues, asset) {
 		st.CurrentBalanceQuery[account] = append(previousValues, asset)
 	}
 }

internal/interpreter/interpreter.go

@@ -277,6 +277,10 @@ func (st *programState) runStatement(statement parser.Statement) ([]Posting, Int
 
 	case *parser.SendStatement:
 		return st.runSendStatement(*statement)
+
+	case *parser.SaveStatement:
+		return st.runSaveStatement(*statement)
+
 	default:
 		utils.NonExhaustiveMatchPanic[any](statement)
 		return nil, nil
@@ -299,6 +303,41 @@ func (st *programState) getPostings() ([]Posting, InterpreterError) {
 	return postings, nil
 }
 
+func (st *programState) runSaveStatement(saveStatement parser.SaveStatement) ([]Posting, InterpreterError) {
+	asset, amt, err := st.evaluateSentAmt(saveStatement.SentValue)
+	if err != nil {
+		return nil, err
+	}
+
+	account, err := evaluateLitExpecting(st, saveStatement.Literal, expectAccount)
+	if err != nil {
+		return nil, err
+	}
+
+	balance := st.getCachedBalance(*account, *asset)
+
+	if amt == nil {
+		balance.Set(big.NewInt(0))
+	} else {
+		// Do not allow negative saves
+		if amt.Cmp(big.NewInt(0)) == -1 {
+			return nil, NegativeAmountErr{
+				Range:  saveStatement.SentValue.GetRange(),
+				Amount: MonetaryInt(*amt),
+			}
+		}
+
+		// we decrease the balance by "amt"
+		balance.Sub(balance, amt)
+		// without going under 0
+		if balance.Cmp(big.NewInt(0)) == -1 {
+			balance.Set(big.NewInt(0))
+		}
+	}
+
+	return nil, nil
+}
+
 func (st *programState) runSendStatement(statement parser.SendStatement) ([]Posting, InterpreterError) {
 	switch sentValue := statement.SentValue.(type) {
 	case *parser.SentValueAll:
@@ -796,3 +835,27 @@ func setAccountMeta(st *programState, r parser.Range, args []Value) InterpreterE
 
 	return nil
 }
+
+func (st *programState) evaluateSentAmt(sentValue parser.SentValue) (*string, *big.Int, InterpreterError) {
+	switch sentValue := sentValue.(type) {
+	case *parser.SentValueAll:
+		asset, err := evaluateLitExpecting(st, sentValue.Asset, expectAsset)
+		if err != nil {
+			return nil, nil, err
+		}
+		return asset, nil, nil
+
+	case *parser.SentValueLiteral:
+		monetary, err := evaluateLitExpecting(st, sentValue.Monetary, expectMonetary)
+		if err != nil {
+			return nil, nil, err
+		}
+		s := string(monetary.Asset)
+		bi := big.Int(monetary.Amount)
+		return &s, &bi, nil
+
+	default:
+		utils.NonExhaustiveMatchPanic[any](sentValue)
+		return nil, nil, nil
+	}
+}