Fix colored asset bug

internal/interpreter/balances.go

@@ -43,10 +43,10 @@ func coloredAsset(asset string, color *string) string {
 
 // 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 {
+func (b Balances) fetchBalance(account string, uncoloredAsset string, color string) *big.Int {
 	accountBalances := b.fetchAccountBalances(account)
 
-	return defaultMapGet(accountBalances, asset, func() *big.Int {
+	return defaultMapGet(accountBalances, coloredAsset(uncoloredAsset, &color), func() *big.Int {
 		return new(big.Int)
 	})
 }

internal/interpreter/interpreter.go

@@ -313,7 +313,7 @@ func (st *programState) pushSender(name string, monetary *big.Int, color string)
 		return
 	}
 
-	balance := st.CachedBalances.fetchBalance(name, st.CurrentAsset)
+	balance := st.CachedBalances.fetchBalance(name, st.CurrentAsset, color)
 	balance.Sub(balance, monetary)
 
 	st.fundsStack.Push(Sender{Name: name, Amount: monetary, Color: color})
@@ -336,13 +336,13 @@ func (st *programState) pushReceiver(name string, monetary *big.Int) {
 
 		if name == KEPT_ADDR {
 			// If funds are kept, give them back to senders
-			srcBalance := st.CachedBalances.fetchBalance(postings.Source, postings.Asset)
+			srcBalance := st.CachedBalances.fetchBalance(postings.Source, st.CurrentAsset, sender.Color)
 			srcBalance.Add(srcBalance, postings.Amount)
 
 			continue
 		}
 
-		destBalance := st.CachedBalances.fetchBalance(postings.Destination, postings.Asset)
+		destBalance := st.CachedBalances.fetchBalance(postings.Destination, st.CurrentAsset, sender.Color)
 		destBalance.Add(destBalance, postings.Amount)
 
 		st.Postings = append(st.Postings, postings)
@@ -389,7 +389,7 @@ func (st *programState) runSaveStatement(saveStatement parser.SaveStatement) Int
 		return err
 	}
 
-	balance := st.CachedBalances.fetchBalance(*account, *asset)
+	balance := st.CachedBalances.fetchBalance(*account, *asset, "")
 
 	if amt == nil {
 		balance.Set(big.NewInt(0))
@@ -478,7 +478,7 @@ func (s *programState) sendAllToAccount(accountLiteral parser.ValueExpr, overdra
 		return nil, err
 	}
 
-	balance := s.CachedBalances.fetchBalance(*account, coloredAsset(s.CurrentAsset, color))
+	balance := s.CachedBalances.fetchBalance(*account, s.CurrentAsset, *color)
 
 	// we sent balance+overdraft
 	sentAmt := CalculateMaxSafeWithdraw(balance, overdraft)
@@ -588,7 +588,7 @@ func (s *programState) trySendingToAccount(accountLiteral parser.ValueExpr, amou
 		// unbounded overdraft: we send the required amount
 		actuallySentAmt = new(big.Int).Set(amount)
 	} else {
-		balance := s.CachedBalances.fetchBalance(*account, coloredAsset(s.CurrentAsset, color))
+		balance := s.CachedBalances.fetchBalance(*account, s.CurrentAsset, *color)
 
 		// that's the amount we are allowed to send (balance + overdraft)
 		actuallySentAmt = CalculateSafeWithdraw(balance, overdraft, amount)
@@ -885,7 +885,7 @@ func getBalance(
 	if fetchBalanceErr != nil {
 		return nil, QueryBalanceError{WrappedError: fetchBalanceErr}
 	}
-	balance := s.CachedBalances.fetchBalance(account, asset)
+	balance := s.CachedBalances.fetchBalance(account, asset, "")
 	return balance, nil
 
 }

internal/interpreter/interpreter_test.go

@@ -4362,6 +4362,50 @@ func TestColorInorder(t *testing.T) {
 	testWithFeatureFlag(t, tc, flags.ExperimentalAssetColors)
 }
 
+func TestColorInorderSendAll(t *testing.T) {
+
+	script := `
+ 		send [COIN *] (
+			source = {
+					@src \ "RED"
+					@src \ "BLUE"
+					@src
+			}
+			destination = @dest
+		)
+	`
+
+	tc := NewTestCase()
+	tc.setBalance("src", "COIN", 100)
+	tc.setBalance("src", "COIN_RED", 20)
+	tc.setBalance("src", "COIN_BLUE", 30)
+	tc.compile(t, script)
+
+	tc.expected = CaseResult{
+		Postings: []Posting{
+			{
+				Asset:       "COIN_RED",
+				Amount:      big.NewInt(20),
+				Source:      "src",
+				Destination: "dest",
+			},
+			{
+				Asset:       "COIN_BLUE",
+				Amount:      big.NewInt(30),
+				Source:      "src",
+				Destination: "dest",
+			},
+			{
+				Asset:       "COIN",
+				Amount:      big.NewInt(100),
+				Source:      "src",
+				Destination: "dest",
+			},
+		},
+	}
+	testWithFeatureFlag(t, tc, flags.ExperimentalAssetColors)
+}
+
 func TestEmptyColor(t *testing.T) {
 	// empty string color behaves as no color