Fix tracked MemoryBlocks

src/main/java/com/jozufozu/flywheel/api/struct/StructWriter.java

@@ -3,11 +3,11 @@ package com.jozufozu.flywheel.api.struct;
 import com.jozufozu.flywheel.api.instancer.InstancedPart;
 
 /**
- * StructWriters can quickly consume many instances of S and write them to some backing buffer.
+ * StructWriters can quickly consume many instances of S and write them to some memory address.
  */
 public interface StructWriter<S extends InstancedPart> {
 	/**
-	 * Write the given struct to given memory address.
+	 * Write the given struct to the given memory address.
 	 */
 	void write(long ptr, S struct);
 }

src/main/java/com/jozufozu/flywheel/backend/memory/FlwMemoryTracker.java

@@ -21,50 +21,40 @@ public class FlwMemoryTracker {
 
 	public static MemoryBlock mallocBlock(long size) {
 		MemoryBlock block = new MemoryBlockImpl(malloc(size), size);
-		cpuMemory += block.size();
+		_allocCPUMemory(block.size());
 		return block;
 	}
 
 	public static MemoryBlock mallocBlockTracked(long size) {
-		TrackedMemoryBlockImpl block = new TrackedMemoryBlockImpl(malloc(size), size);
-		block.cleanable = CLEANER.register(block, () -> {
-			if (!block.isFreed()) {
-				block.free();
-			}
-		});
-		cpuMemory += block.size();
+		MemoryBlock block = new TrackedMemoryBlockImpl(malloc(size), size, CLEANER);
+		_allocCPUMemory(block.size());
 		return block;
 	}
 
 	@Deprecated
 	public static ByteBuffer mallocBuffer(int size) {
 		ByteBuffer buffer = MemoryUtil.memByteBuffer(malloc(size), size);
-		cpuMemory += buffer.capacity();
+		_allocCPUMemory(buffer.capacity());
 		return buffer;
 	}
 
 	public static long calloc(long num, long size) {
 		long ptr = MemoryUtil.nmemCalloc(num, size);
 		if (ptr == MemoryUtil.NULL) {
-			throw new OutOfMemoryError("Failed to allocate " + num + " blocks of size " + size + " bytes");
+			throw new OutOfMemoryError("Failed to allocate " + num + " elements of size " + size + " bytes");
 		}
 		return ptr;
 	}
 
 	public static MemoryBlock callocBlock(long num, long size) {
 		MemoryBlock block = new MemoryBlockImpl(calloc(num, size), num * size);
-		cpuMemory += block.size();
+		_allocCPUMemory(block.size());
 		return block;
 	}
 
 	public static MemoryBlock callocBlockTracked(long num, long size) {
-		TrackedMemoryBlockImpl block = new TrackedMemoryBlockImpl(calloc(num, size), num * size);
-		block.cleanable = CLEANER.register(block, () -> {
-			if (!block.isFreed()) {
-				block.free();
-			}
-		});
-		cpuMemory += block.size();
+		MemoryBlock block = new TrackedMemoryBlockImpl(calloc(num, size), num * size, CLEANER);
+		_allocCPUMemory(block.size());
 		return block;
 	}
 
@@ -78,25 +68,23 @@ public class FlwMemoryTracker {
 
 	public static MemoryBlock reallocBlock(MemoryBlock block, long size) {
 		MemoryBlock newBlock = new MemoryBlockImpl(realloc(block.ptr(), size), size);
-		cpuMemory += -block.size() + newBlock.size();
+		_freeCPUMemory(block.size());
+		_allocCPUMemory(newBlock.size());
 		return newBlock;
 	}
 
 	public static MemoryBlock reallocBlockTracked(MemoryBlock block, long size) {
-		TrackedMemoryBlockImpl newBlock = new TrackedMemoryBlockImpl(realloc(block.ptr(), size), size);
-		newBlock.cleanable = CLEANER.register(newBlock, () -> {
-			if (!newBlock.isFreed()) {
-				newBlock.free();
-			}
-		});
-		cpuMemory += -block.size() + newBlock.size();
+		MemoryBlock newBlock = new TrackedMemoryBlockImpl(realloc(block.ptr(), size), size, CLEANER);
+		_freeCPUMemory(block.size());
+		_allocCPUMemory(newBlock.size());
 		return newBlock;
 	}
 
 	@Deprecated
 	public static ByteBuffer reallocBuffer(ByteBuffer buffer, int size) {
 		ByteBuffer newBuffer = MemoryUtil.memByteBuffer(realloc(MemoryUtil.memAddress(buffer), size), size);
-		cpuMemory += -buffer.capacity() + newBuffer.capacity();
+		_freeCPUMemory(buffer.capacity());
+		_allocCPUMemory(newBuffer.capacity());
 		return newBuffer;
 	}
 
@@ -106,13 +94,13 @@ public class FlwMemoryTracker {
 
 	public static void freeBlock(MemoryBlock block) {
 		free(block.ptr());
-		cpuMemory -= block.size();
+		_freeCPUMemory(block.size());
 	}
 
 	@Deprecated
 	public static void freeBuffer(ByteBuffer buffer) {
 		free(MemoryUtil.memAddress(buffer));
-		cpuMemory -= buffer.capacity();
+		_freeCPUMemory(buffer.capacity());
 	}
 
 	public static void _allocCPUMemory(long size) {

src/main/java/com/jozufozu/flywheel/backend/memory/MemoryBlock.java

@@ -13,9 +13,7 @@ public sealed interface MemoryBlock permits MemoryBlockImpl {
 
 	void copyTo(long ptr, long bytes);
 
-	default void copyTo(long ptr) {
-		copyTo(ptr, size());
-	}
+	void copyTo(long ptr);
 
 	void clear();
 

src/main/java/com/jozufozu/flywheel/backend/memory/MemoryBlockImpl.java

@@ -40,6 +40,11 @@ sealed class MemoryBlockImpl implements MemoryBlock permits TrackedMemoryBlockIm
 		MemoryUtil.memCopy(this.ptr, ptr, bytes);
 	}
 
+	@Override
+	public void copyTo(long ptr) {
+		copyTo(ptr, size);
+	}
+
 	@Override
 	public void clear() {
 		MemoryUtil.memSet(ptr, 0, size);

src/main/java/com/jozufozu/flywheel/backend/memory/TrackedMemoryBlockImpl.java

@@ -3,10 +3,13 @@ package com.jozufozu.flywheel.backend.memory;
 import java.lang.ref.Cleaner;
 
 final class TrackedMemoryBlockImpl extends MemoryBlockImpl {
-	Cleaner.Cleanable cleanable;
+	final CleaningAction cleaningAction;
+	final Cleaner.Cleanable cleanable;
 
-	TrackedMemoryBlockImpl(long ptr, long size) {
+	TrackedMemoryBlockImpl(long ptr, long size, Cleaner cleaner) {
 		super(ptr, size);
+		cleaningAction = new CleaningAction(ptr, size);
+		cleanable = cleaner.register(this, cleaningAction);
 	}
 
 	@Override
@@ -14,25 +17,49 @@ final class TrackedMemoryBlockImpl extends MemoryBlockImpl {
 		return true;
 	}
 
+	void freeInner() {
+		freed = true;
+		cleaningAction.freed = true;
+		cleanable.clean();
+	}
+
 	@Override
 	public MemoryBlock realloc(long size) {
 		MemoryBlock block = FlwMemoryTracker.reallocBlock(this, size);
-		freed = true;
-		cleanable.clean();
+		freeInner();
 		return block;
 	}
 
 	@Override
 	public MemoryBlock reallocTracked(long size) {
 		MemoryBlock block = FlwMemoryTracker.reallocBlockTracked(this, size);
-		freed = true;
-		cleanable.clean();
+		freeInner();
 		return block;
 	}
 
 	@Override
 	public void free() {
-		cleanable.clean();
-		freed = true;
+		FlwMemoryTracker.freeBlock(this);
+		freeInner();
+	}
+
+	static class CleaningAction implements Runnable {
+		final long ptr;
+		final long size;
+
+		boolean freed;
+
+		CleaningAction(long ptr, long size) {
+			this.ptr = ptr;
+			this.size = size;
+		}
+
+		@Override
+		public void run() {
+			if (!freed) {
+				FlwMemoryTracker.free(ptr);
+				FlwMemoryTracker._freeCPUMemory(size);
+			}
+		}
 	}
 }