From 7c2662d7fe260daa0298386af1c7a806371ac794 Mon Sep 17 00:00:00 2001 From: munna0908 Date: Thu, 10 Jul 2025 18:55:43 +0530 Subject: [PATCH] use uniqueptr for safe memory managment --- codex/merkletree/codex/codex.nim | 5 +- codex/merkletree/merkletree.nim | 76 ++------------------ codex/merkletree/poseidon2.nim | 4 +- codex/nat.nim | 2 +- codex/utils/poseidon2digest.nim | 19 ++--- codex/utils/uniqueptr.nim | 58 +++++++++++++++ tests/codex/merkletree/testposeidon2tree.nim | 2 - 7 files changed, 74 insertions(+), 92 deletions(-) create mode 100644 codex/utils/uniqueptr.nim diff --git a/codex/merkletree/codex/codex.nim b/codex/merkletree/codex/codex.nim index 67e1507e..dc4544c0 100644 --- a/codex/merkletree/codex/codex.nim +++ b/codex/merkletree/codex/codex.nim @@ -25,6 +25,8 @@ import ../../blocktype from ../../utils/digest import digestBytes +import ../../utils/uniqueptr + import ../merkletree export merkletree @@ -208,9 +210,6 @@ proc init*( if not task.success.load(): return failure("merkle tree task failed") - # defer: - # task.layers = default(Isolated[seq[seq[ByteHash]]]) - tree.layers = extractValue(task.layers) success tree diff --git a/codex/merkletree/merkletree.nim b/codex/merkletree/merkletree.nim index bf9d4c92..5c12f198 100644 --- a/codex/merkletree/merkletree.nim +++ b/codex/merkletree/merkletree.nim @@ -16,68 +16,7 @@ import pkg/taskpools import pkg/chronos/threadsync import ../errors - -type UniqueSeq*[T] = object - ## A unique pointer to a seq[seq[T]] in shared memory - ## Can only be moved, not copied - data: ptr seq[seq[T]] - -proc newUniqueSeq*[T](data: sink Isolated[seq[seq[T]]]): UniqueSeq[T] = - ## Creates a new unique sequence in shared memory - ## The memory is automatically freed when the object is destroyed - result.data = cast[ptr seq[seq[T]]](allocShared0(sizeof(seq[seq[T]]))) - - result.data[] = extract(data) - -proc `=destroy`*[T](p: var UniqueSeq[T]) = - ## Destructor for UniqueSeq - if p.data != nil: - # Clear the sequence to release inner sequences - p.data[].setLen(0) - echo "destroying unique seq" - deallocShared(p.data) - p.data = nil - -proc `=copy`*[T]( - dest: var UniqueSeq[T], src: UniqueSeq[T] -) {.error: "UniqueSeq cannot be copied, only moved".} - -proc `=sink`*[T](dest: var UniqueSeq[T], src: UniqueSeq[T]) = - ## Move constructor for UniqueSeq - if dest.data != nil: - `=destroy`(dest) - dest.data = src.data - # We need to nil out the source data to prevent double-free - # This is handled by Nim's destructive move semantics - -proc `[]`*[T](p: UniqueSeq[T]): lent seq[seq[T]] = - ## Access the data (read-only) - if p.data == nil: - raise newException(NilAccessDefect, "accessing nil UniqueSeq") - p.data[] - -proc `[]`*[T](p: var UniqueSeq[T]): var seq[seq[T]] = - ## Access the data (mutable) - if p.data == nil: - raise newException(NilAccessDefect, "accessing nil UniqueSeq") - p.data[] - -proc isNil*[T](p: UniqueSeq[T]): bool = - ## Check if the UniqueSeq is nil - p.data == nil - -proc extractValue*[T](p: var UniqueSeq[T]): seq[seq[T]] = - ## Extract the value from the UniqueSeq and release the memory - if p.data == nil: - raise newException(NilAccessDefect, "extracting from nil UniqueSeq") - - # Move the value out - var isolated = isolate(p.data[]) - result = extract(isolated) - - # Free the shared memory - deallocShared(p.data) - p.data = nil +import ../utils/uniqueptr type CompressFn*[H, K] = proc(x, y: H, key: K): ?!H {.noSideEffect, raises: [].} @@ -98,7 +37,7 @@ type tree*: ptr MerkleTree[H, K] leaves*: seq[H] signal*: ThreadSignalPtr - layers*: UniqueSeq[H] + layers*: UniquePtr[seq[seq[H]]] success*: Atomic[bool] func depth*[H, K](self: MerkleTree[H, K]): int = @@ -233,12 +172,11 @@ proc merkleTreeWorker*[H, K](task: ptr MerkleTask[H, K]) {.gcsafe.} = task[].success.store(false) return - var l = res.get() - var newOuterSeq = newSeq[seq[H]](l.len) - for i in 0 ..< l.len: - var isoInner = isolate(l[i]) + var layers = res.get() + var newOuterSeq = newSeq[seq[H]](layers.len) + for i in 0 ..< layers.len: + var isoInner = isolate(layers[i]) newOuterSeq[i] = extract(isoInner) - var isolatedLayers = isolate(newOuterSeq) - task[].layers = newUniqueSeq(isolatedLayers) + task[].layers = newUniquePtr(newOuterSeq) task[].success.store(true) diff --git a/codex/merkletree/poseidon2.nim b/codex/merkletree/poseidon2.nim index 0917276c..6feb9df4 100644 --- a/codex/merkletree/poseidon2.nim +++ b/codex/merkletree/poseidon2.nim @@ -19,6 +19,7 @@ import pkg/constantine/platforms/abstractions import pkg/questionable/results import ../utils +import ../utils/uniqueptr import ../rng import ./merkletree @@ -118,9 +119,6 @@ proc init*( if not task.success.load(): return failure("merkle tree task failed") - # defer: - # task.layers = default(Isolated[seq[seq[Poseidon2Hash]]]) - tree.layers = extractValue(task.layers) success tree diff --git a/codex/nat.nim b/codex/nat.nim index d022dad6..275367df 100644 --- a/codex/nat.nim +++ b/codex/nat.nim @@ -423,7 +423,7 @@ proc nattedAddress*( it.remapAddr(ip = newIP, port = tcp) else: # NAT mapping failed - use original address - echo "Failed to get external IP, using original address", it + error "Failed to get external IP, using original address", it discoveryAddrs.add(getMultiAddrWithIPAndUDPPort(ipPart.get, udpPort)) it else: diff --git a/codex/utils/poseidon2digest.nim b/codex/utils/poseidon2digest.nim index 7e7c37f2..c023928e 100644 --- a/codex/utils/poseidon2digest.nim +++ b/codex/utils/poseidon2digest.nim @@ -15,6 +15,7 @@ import pkg/stew/byteutils import pkg/taskpools import pkg/chronos import pkg/chronos/threadsync +import ./uniqueptr import ../merkletree @@ -23,7 +24,7 @@ type DigestTask* = object bytes: seq[byte] chunkSize: int success: Atomic[bool] - digest: ptr Poseidon2Hash + digest: UniquePtr[Poseidon2Hash] export DigestTask @@ -85,8 +86,7 @@ proc digestWorker(tp: Taskpool, task: ptr DigestTask) {.gcsafe.} = task[].success.store(false) return - var isolatedDigest = isolate(res.get()) - task[].digest[] = extract(isolatedDigest) + task[].digest = newUniquePtr(res.get()) task[].success.store(true) proc digest*( @@ -100,12 +100,7 @@ proc digest*( doAssert tp.numThreads > 1, "Must have at least one separate thread or signal will never be fired" - var task = DigestTask( - signal: signal, - bytes: bytes, - chunkSize: chunkSize, - digest: cast[ptr Poseidon2Hash](allocShared(sizeof(Poseidon2Hash))), - ) + var task = DigestTask(signal: signal, bytes: bytes, chunkSize: chunkSize) tp.spawn digestWorker(tp, addr task) @@ -119,11 +114,7 @@ proc digest*( if not task.success.load(): return failure("digest task failed") - var isolatedDigest = isolate(task.digest[]) - var digest = extract(isolatedDigest) - defer: - deallocShared(task.digest) - success digest + success extractValue(task.digest) func digestMhash*( _: type Poseidon2Tree, bytes: openArray[byte], chunkSize: int diff --git a/codex/utils/uniqueptr.nim b/codex/utils/uniqueptr.nim new file mode 100644 index 00000000..2aec0d38 --- /dev/null +++ b/codex/utils/uniqueptr.nim @@ -0,0 +1,58 @@ +import std/isolation +type UniquePtr*[T] = object + ## A unique pointer to a seq[seq[T]] in shared memory + ## Can only be moved, not copied + data: ptr T + +proc newUniquePtr*[T](data: sink Isolated[T]): UniquePtr[T] = + ## Creates a new unique sequence in shared memory + ## The memory is automatically freed when the object is destroyed + result.data = cast[ptr T](allocShared0(sizeof(T))) + result.data[] = extract(data) + +template newUniquePtr*[T](data: T): UniquePtr[T] = + newUniquePtr(isolate(data)) + +proc `=destroy`*[T](p: var UniquePtr[T]) = + ## Destructor for UniquePtr + if p.data != nil: + deallocShared(p.data) + p.data = nil + +proc `=copy`*[T]( + dest: var UniquePtr[T], src: UniquePtr[T] +) {.error: "UniquePtr cannot be copied, only moved".} + +proc `=sink`*[T](dest: var UniquePtr[T], src: UniquePtr[T]) = + if dest.data != nil: + `=destroy`(dest) + dest.data = src.data + # We need to nil out the source data to prevent double-free + # This is handled by Nim's destructive move semantics + +proc `[]`*[T](p: UniquePtr[T]): lent T = + ## Access the data (read-only) + if p.data == nil: + raise newException(NilAccessDefect, "accessing nil UniquePtr") + p.data[] + +# proc `[]`*[T](p: var UniquePtr[T]): var T = +# ## Access the data (mutable) +# if p.data == nil: +# raise newException(NilAccessDefect, "accessing nil UniquePtr") +# p.data[] + +proc isNil*[T](p: UniquePtr[T]): bool = + ## Check if the UniquePtr is nil + p.data == nil + +proc extractValue*[T](p: var UniquePtr[T]): T = + ## Extract the value from the UniquePtr and release the memory + if p.data == nil: + raise newException(NilAccessDefect, "extracting from nil UniquePtr") + # Move the value out + var isolated = isolate(p.data[]) + result = extract(isolated) + # Free the shared memory + deallocShared(p.data) + p.data = nil diff --git a/tests/codex/merkletree/testposeidon2tree.nim b/tests/codex/merkletree/testposeidon2tree.nim index f574d637..45a727e5 100644 --- a/tests/codex/merkletree/testposeidon2tree.nim +++ b/tests/codex/merkletree/testposeidon2tree.nim @@ -63,11 +63,9 @@ suite "Test Poseidon2Tree": tree == fromNodes test "Build poseidon2 tree from poseidon2 leaves asynchronously": - echo "Build poseidon2 tree from poseidon2 leaves asynchronously" var tp = Taskpool.new() defer: tp.shutdown() - echo "@@@@@" let tree = (await Poseidon2Tree.init(tp, leaves = expectedLeaves)).tryGet() check: