From f2f1b41e87fcd5c5925ffae81f386823e1b99202 Mon Sep 17 00:00:00 2001
From: Artem <aopoltorzhicky@gmail.com>
Date: Sun, 2 Jun 2024 12:29:35 +0200
Subject: [PATCH] Fix: lazy storage diff parsing

---
 internal/parsers/operations/transaction.go |  2 +-
 internal/parsers/protocols/migration.go    |  2 +-
 internal/parsers/storage/lazy_babylon.go   | 22 ++++++++++++++++------
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/internal/parsers/operations/transaction.go b/internal/parsers/operations/transaction.go
index d451b5879..70578a09b 100644
--- a/internal/parsers/operations/transaction.go
+++ b/internal/parsers/operations/transaction.go
@@ -93,7 +93,7 @@ func (p Transaction) Parse(ctx context.Context, data noderpc.Operation, store pa
 	}
 }
 
-func (p Transaction) parseSmartRollupParams(data noderpc.Operation, tx *operation.Operation) error {
+func (p Transaction) parseSmartRollupParams(_ noderpc.Operation, tx *operation.Operation) error {
 	if len(tx.Parameters) == 0 {
 		return tx.Entrypoint.Set(consts.DefaultEntrypoint)
 	}
diff --git a/internal/parsers/protocols/migration.go b/internal/parsers/protocols/migration.go
index d598d8211..51b68d1e5 100644
--- a/internal/parsers/protocols/migration.go
+++ b/internal/parsers/protocols/migration.go
@@ -112,7 +112,7 @@ func (m Migration) contractMigrations(
 	return nil
 }
 
-func (m Migration) vestingMigration(ctx context.Context, tx models.Transaction, head noderpc.Header, currentProtocol protocol.Protocol) error {
+func (m Migration) vestingMigration(ctx context.Context, _ models.Transaction, head noderpc.Header, currentProtocol protocol.Protocol) error {
 	addresses, err := m.ctx.RPC.GetContractsByBlock(ctx, head.Level)
 	if err != nil {
 		return err
diff --git a/internal/parsers/storage/lazy_babylon.go b/internal/parsers/storage/lazy_babylon.go
index 9380bca80..0f2c3c773 100644
--- a/internal/parsers/storage/lazy_babylon.go
+++ b/internal/parsers/storage/lazy_babylon.go
@@ -22,6 +22,7 @@ type LazyBabylon struct {
 
 	ptrMap            map[int64]int64
 	temporaryPointers map[int64]*ast.BigMap
+	allocated         map[int64]struct{}
 }
 
 // NewLazyBabylon -
@@ -33,6 +34,7 @@ func NewLazyBabylon(repo bigmapdiff.Repository, operations operation.Repository,
 
 		ptrMap:            make(map[int64]int64),
 		temporaryPointers: make(map[int64]*ast.BigMap),
+		allocated:         make(map[int64]struct{}),
 	}
 }
 
@@ -123,6 +125,9 @@ func (b *LazyBabylon) checkPointers(result *noderpc.OperationResult, storage *as
 			if ptr < 0 {
 				continue
 			}
+			if _, ok := b.temporaryPointers[ptr]; ok {
+				continue
+			}
 
 		default:
 			ptr := lsd.ID
@@ -135,6 +140,9 @@ func (b *LazyBabylon) checkPointers(result *noderpc.OperationResult, storage *as
 			if ptr < 0 {
 				continue
 			}
+			if _, ok := b.temporaryPointers[ptr]; ok {
+				continue
+			}
 
 		}
 		return ErrUnknownTemporaryPointer
@@ -281,6 +289,7 @@ func (b *LazyBabylon) handleBigMapDiffAlloc(ctx context.Context, diff *noderpc.L
 	if ptr > -1 {
 		operation.BigMapActions = append(operation.BigMapActions, b.createBigMapDiffAction("alloc", address, &ptr, nil, operation))
 	}
+	b.allocated[ptr] = struct{}{}
 
 	return b.handleBigMapDiffUpdate(ctx, diff, ptr, address, operation, store)
 }
@@ -347,19 +356,20 @@ func (b *LazyBabylon) updateTemporaryPointers(src, dst int64) {
 
 func (b *LazyBabylon) getCopyBigMapDiff(ctx context.Context, src int64, address string) (bmd []bigmapdiff.BigMapDiff, err error) {
 	if src > -1 {
+		if _, ok := b.allocated[src]; ok {
+			return b.getDiffsFromUpdates(src)
+		}
 		states, err := b.repo.GetByPtr(ctx, address, src)
 		if err != nil {
 			return nil, err
 		}
+
 		bmd = make([]bigmapdiff.BigMapDiff, 0, len(states))
 		for i := range states {
 			bmd = append(bmd, states[i].ToDiff())
 		}
-	} else {
-		bmd, err = b.getDiffsFromUpdates(src)
-		if err != nil {
-			return nil, err
-		}
+		return bmd, nil
 	}
-	return
+
+	return b.getDiffsFromUpdates(src)
 }