Add MemorySegmentManager.GetFeltRange (#272)

* Add MemorySegmentManager.GetFeltRange * Add unit test + comment * Update segments.go --------- Co-authored-by: Mariano A. Nicolini <[email protected]>
lambdaclass · Sep 20, 2023 · 2fa2e5f · 2fa2e5f
1 parent abafc37
commit 2fa2e5f
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 0 deletions.
diff --git a/pkg/vm/memory/segments.go b/pkg/vm/memory/segments.go
@@ -172,3 +172,17 @@ func (m *MemorySegmentManager) Finalize(size *uint, segmentIndex uint, publicMem
 		m.PublicMemoryOffsets[segmentIndex] = emptyList
 	}
 }
+
+// Gets a range of Felt memory values from addr to addr + size
+// Fails if any of the values inside the range is missing (memory gap), or is not a Felt
+func (m *MemorySegmentManager) GetFeltRange(start Relocatable, size uint) ([]lambdaworks.Felt, error) {
+	feltRange := make([]lambdaworks.Felt, 0, size)
+	for i := uint(0); i < size; i++ {
+		val, err := m.Memory.GetFelt(start.AddUint(i))
+		if err != nil {
+			return nil, err
+		}
+		feltRange = append(feltRange, val)
+	}
+	return feltRange, nil
+}
diff --git a/pkg/vm/memory/segments_test.go b/pkg/vm/memory/segments_test.go
@@ -250,3 +250,63 @@ func TestGetMemoryHoles(t *testing.T) {
 		t.Errorf("Get Memory Holes Returned the wrong value. Expected: 2, got %d", result)
 	}
 }
+
+func TestGetFeltRangeOk(t *testing.T) {
+	segments := memory.NewMemorySegmentManager()
+	segments.AddSegment()
+	segments.Memory.Insert(memory.NewRelocatable(0, 1), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(1)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 2), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(2)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 3), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(3)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 4), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(4)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 5), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(5)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 6), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(6)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 7), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(7)))
+
+	feltRange, err := segments.GetFeltRange(memory.NewRelocatable(0, 2), 4)
+	expectedFeltRange := []lambdaworks.Felt{
+		lambdaworks.FeltFromUint64(2),
+		lambdaworks.FeltFromUint64(3),
+		lambdaworks.FeltFromUint64(4),
+		lambdaworks.FeltFromUint64(5),
+	}
+
+	if err != nil || !reflect.DeepEqual(feltRange, expectedFeltRange) {
+		t.Errorf("GetFeltRange failed or returned wrong value")
+	}
+}
+
+func TestGetFeltRangeGap(t *testing.T) {
+	segments := memory.NewMemorySegmentManager()
+	segments.AddSegment()
+	segments.Memory.Insert(memory.NewRelocatable(0, 1), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(1)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 2), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(2)))
+	//segments.Memory.Insert(memory.NewRelocatable(0, 3), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(3)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 4), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(4)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 5), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(5)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 6), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(6)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 7), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(7)))
+
+	_, err := segments.GetFeltRange(memory.NewRelocatable(0, 2), 4)
+
+	if err == nil {
+		t.Errorf("GetFeltRange should have failed")
+	}
+}
+
+func TestGetFeltRangeRelocatable(t *testing.T) {
+	segments := memory.NewMemorySegmentManager()
+	segments.AddSegment()
+	segments.Memory.Insert(memory.NewRelocatable(0, 1), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(1)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 2), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(2)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 3), memory.NewMaybeRelocatableRelocatable(memory.NewRelocatable(0, 0)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 4), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(4)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 5), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(5)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 6), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(6)))
+	segments.Memory.Insert(memory.NewRelocatable(0, 7), memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint64(7)))
+
+	_, err := segments.GetFeltRange(memory.NewRelocatable(0, 2), 4)
+
+	if err == nil {
+		t.Errorf("GetFeltRange should have failed")
+	}
+}