use lang:bs;
use lang:bs:macro;
use core:lang;
use core:asm;

class MapExpr extends Block {
	private Expr src;
	private Expr? transform;
	private LocalVar boundVar;

	init(SrcPos pos, Block parent, Expr src, SStr varName) {
		init(pos, parent) {
			src = expectCastTo(src, named{Array<Int>}, parent.scope);
			boundVar = LocalVar(varName.v, named{Int}, varName.pos, false);
		}

		add(boundVar);
	}

	void setTransform(Expr expr) {
		transform = expectCastTo(expr, named{Int}, scope);
	}

	ExprResult result() : override {
		ExprResult(Value(named{Array<Int>}));
	}

	void blockCode(CodeGen gen, CodeResult res) : override {
		unless (transform)
			throw InternalError("No transform set for a MapBlock!");

		// Create the array. Put it in the result variable:
		var resultVar = res.location(gen);
		allocObject(gen, named{Array<Int>:__init<Array<Int>>}, [], resultVar);
		res.created(gen);

		// Get the original array:
		CodeResult originalArrayRes(named{Array<Int>}, gen.block);
		src.code(gen, originalArrayRes);
		var originalArray = originalArrayRes.location(gen);

		// Extract the count of the original array.
		Function countFn = named{Array<Int>:count<Array<Int>>};
		CodeResult originalCountRes(named{Nat}, gen.block);
		countFn.autoCall(gen, [originalArray], originalCountRes);
		var originalCount = originalCountRes.location(gen);

		// Create a loop counter:
		var loopCounter = gen.createVar(named{Nat}).v;
		gen.l << mov(loopCounter, natConst(0));

		// Start the loop:
		Label loopStart = gen.l.label();
		gen.l << loopStart;

		// Check if we are done.
		Label loopEnd = gen.l.label();
		gen.l << cmp(loopCounter, originalCount);
		gen.l << jmp(loopEnd, CondFlag:ifAboveEqual);

		// Get an element. Returns a reference to the element.
		Function accessFunction = named{Array<Int>:"[]"<Array<Int>, Nat>};
		CodeResult elemPtr(accessFunction.result, gen.block);
		accessFunction.autoCall(gen, [originalArray, loopCounter], elemPtr);

		// Evaluate the expression to get the value to insert. First set our variable:
		gen.l << mov(ptrA, elemPtr.location(gen));
		gen.l << mov(boundVar.var.v, intRel(ptrA));
		CodeResult transformed(named{Int}, gen.block);
		transform.code(gen, transformed);

		// Push it into the array.
		Function pushFn = named{Array<Int>:push<Array<Int>, Int>};
		gen.l << lea(ptrA, transformed.location(gen));
		pushFn.autoCall(gen, [resultVar, ptrA], CodeResult());

		// Increment the loop counter and restart the loop.
		gen.l << add(loopCounter, natConst(1));
		gen.l << jmp(loopStart);

		// End of the loop.
		gen.l << loopEnd;
	}
}
