From a206589bda0f2564269cbe618da96df1a1632eaa Mon Sep 17 00:00:00 2001
From: Artem <aopoltorzhicky@gmail.com>
Date: Wed, 10 Mar 2021 14:27:39 +0300
Subject: [PATCH] Fix: or path for big map

---
 internal/contractparser/storage/babylon.go | 39 +++++++++++++-------
 internal/contractparser/storage/public.go  | 42 +++++++++++++++++++---
 2 files changed, 64 insertions(+), 17 deletions(-)

diff --git a/internal/contractparser/storage/babylon.go b/internal/contractparser/storage/babylon.go
index f5298b047..bcce5ebb8 100644
--- a/internal/contractparser/storage/babylon.go
+++ b/internal/contractparser/storage/babylon.go
@@ -73,7 +73,7 @@ func (b *Babylon) ParseTransaction(content gjson.Result, metadata meta.Metadata,
 		return RichStorage{Empty: true}, err
 	}
 	if result.Get("big_map_diff.#").Int() > 0 {
-		ptrMap, err := FindBigMapPointers(metadata, normalizedStorage)
+		ptrMap, err := b.findPointers(metadata, normalizedStorage, operation)
 		if err != nil {
 			return RichStorage{Empty: true}, err
 		}
@@ -100,18 +100,9 @@ func (b *Babylon) ParseOrigination(content gjson.Result, metadata meta.Metadata,
 
 	var bm []models.Model
 	if result.Get("big_map_diff.#").Int() > 0 {
-		ptrToBin, err := FindBigMapPointers(metadata, storage)
-		if err != nil || len(ptrToBin) == 0 {
-			// If pointers are not found into script`s storage we try to receive storage from node and find pointers there
-			// If pointers are not found after that -> throw error
-			storage, err = b.rpc.GetScriptStorageJSON(address, operation.Level)
-			if err != nil {
-				return RichStorage{Empty: true}, err
-			}
-			ptrToBin, err = FindBigMapPointers(metadata, storage)
-			if err != nil {
-				return RichStorage{Empty: true}, err
-			}
+		ptrToBin, err := b.findPointers(metadata, storage, operation)
+		if err != nil {
+			return RichStorage{Empty: true}, err
 		}
 
 		if bm, err = b.handleBigMapDiff(result, ptrToBin, address, operation); err != nil {
@@ -204,6 +195,28 @@ func (b *Babylon) Enrich(sStorage, sPrevStorage string, bmd []bigmapdiff.BigMapD
 	return storage, nil
 }
 
+func (b *Babylon) findPointers(metadata meta.Metadata, storage gjson.Result, operation operation.Operation) (map[int64]string, error) {
+	// If pointers are not found into script`s storage we try to receive storage from node by current level and find pointers there
+	// If pointers are not found after that -> we try to receive storage from node by previous level and find pointers there
+	// If pointers are not found after that -> throw error
+	ptrToBin, err := FindBigMapPointers(metadata, storage)
+	if err == nil && len(ptrToBin) > 0 {
+		return ptrToBin, err
+	}
+
+	for i := int64(0); i < 2; i++ {
+		storage, err = b.rpc.GetScriptStorageJSON(operation.Destination, operation.Level-i)
+		if err != nil {
+			return nil, err
+		}
+		ptrToBin, err = FindBigMapPointers(metadata, storage)
+		if err == nil && len(ptrToBin) > 0 {
+			break
+		}
+	}
+	return ptrToBin, err
+}
+
 func (b *Babylon) handleBigMapDiff(result gjson.Result, ptrMap map[int64]string, address string, op operation.Operation) ([]models.Model, error) {
 	storageModels := make([]models.Model, 0)
 
diff --git a/internal/contractparser/storage/public.go b/internal/contractparser/storage/public.go
index 07ef85f8d..c857f7429 100644
--- a/internal/contractparser/storage/public.go
+++ b/internal/contractparser/storage/public.go
@@ -34,6 +34,8 @@ func FindBigMapPointers(m meta.Metadata, storage gjson.Result) (map[int64]string
 			continue
 		}
 
+		k := checkOnOr(k, m)
+
 		if err := setMapPtr(storage, k, key); err != nil {
 			return nil, err
 		}
@@ -41,7 +43,34 @@ func FindBigMapPointers(m meta.Metadata, storage gjson.Result) (map[int64]string
 	return key, nil
 }
 
-func setMapPtr(storage gjson.Result, path string, m map[int64]string) error {
+func checkOnOr(path string, m meta.Metadata) string {
+	var buf strings.Builder
+	parts := strings.Split(path, "/")
+
+	var nextOr bool
+	for i := 0; i < len(parts); i++ {
+		if i > 0 {
+			buf.WriteByte('/')
+		}
+		subPath := strings.Join(parts[:i+1], "/")
+		node, ok := m[subPath]
+		if !ok {
+			return path
+		}
+
+		if nextOr {
+			buf.WriteString("0")
+		} else {
+			buf.WriteString(parts[i])
+		}
+
+		nextOr = node.Prim == consts.OR
+	}
+
+	return buf.String()
+}
+
+func getJSONPath(path string) string {
 	var buf strings.Builder
 
 	trimmed := strings.TrimPrefix(path, "0/")
@@ -61,11 +90,16 @@ func setMapPtr(storage gjson.Result, path string, m map[int64]string) error {
 			buf.WriteString(".")
 		}
 	}
-	buf.WriteString("int")
+	resp := buf.String()
+	return strings.TrimSuffix(resp, ".")
+}
+
+func setMapPtr(storage gjson.Result, path string, m map[int64]string) error {
+	pathJSON := getJSONPath(path) + ".int"
 
-	ptr := storage.Get(buf.String())
+	ptr := storage.Get(pathJSON)
 	if !ptr.Exists() {
-		return errors.Wrapf(ErrPathIsNotPointer, "path=%s buf=%s", path, buf.String())
+		return errors.Wrapf(ErrPathIsNotPointer, "path=%s buf=%s", path, pathJSON)
 	}
 
 	for _, p := range ptr.Array() {