import java.lang.foreign.Arena;
import java.lang.foreign.MemorySegment;
import java.lang.foreign.SegmentAllocator;
import java.lang.foreign.ValueLayout;

// 创建时必须用try-with-resources保证释放.
// 创建出的MemStack对象及其分配出来的MemorySegment只能在当前线程的栈上(包括方法参数)传递,
// 绝对不能传到堆上(对象/类的字段), 以确保生命期在所属的MemStack对象内,且没有并发访问.
public class MemStack implements SegmentAllocator, AutoCloseable {
	private final int blockSize; // 每个内存段的大小
	private final Arena arena = Arena.ofConfined();
	private MemorySegment curSeg; // 当前正在分配的内存段
	private int offset; // curSeg已经分配的偏移
	private int allocSize; // 从arena分配的总大小

	public MemStack() {
		this(4096);
	}

	public MemStack(int blockSize) {
		this.blockSize = blockSize;
	}

	public int getAllocSize() {
		return allocSize;
	}

	@Override
	public MemorySegment allocate(long byteSize, long byteAlignment) {
		if (byteSize < 0)
			throw new IllegalArgumentException("negative byteSize: " + byteSize);
		if (byteAlignment <= 0 || (byteAlignment & (byteAlignment - 1)) != 0)
			throw new IllegalArgumentException("invalid byteAlignment: " + byteAlignment);
		var seg = curSeg;
		if (seg == null) {
			allocSize += blockSize;
			curSeg = seg = arena.allocate(blockSize);
		}
		long off = offset;
		var mask = byteAlignment - 1;
		var lowBits = (seg.address() + off) & mask;
		if (lowBits != 0)
			off += byteAlignment - lowBits;
		var end = off + byteSize;
		if (end <= blockSize) {
			offset = (int)end;
			return seg.asSlice(off, byteSize);
		}
		if (byteSize <= (blockSize >>> 2)) {
			allocSize += blockSize;
			curSeg = seg = arena.allocate(blockSize);
			lowBits = seg.address() & mask;
			off = lowBits != 0 ? byteAlignment - lowBits : 0;
			end = off + byteSize;
			if (end <= blockSize) {
				offset = (int)end;
				return seg.asSlice(off, byteSize);
			}
			offset = 0;
		}
		//noinspection lossy-conversions
		allocSize += byteSize;
		return arena.allocate(byteSize, byteAlignment);
	}

	@Override
	public void close() {
		arena.close();
	}

	public static void main(String[] args) {
		try (var stack = new MemStack()) {
			var v1 = stack.allocateFrom(ValueLayout.JAVA_FLOAT, 12345);
			var v2 = stack.allocateFrom(ValueLayout.JAVA_LONG, 67890);
			System.out.println(stack.curSeg.address());
			System.out.println(v1.address());
			System.out.println(v2.address());
			System.out.println(stack.offset);
			System.out.println(stack.getAllocSize());
		}
	}
}