package org.mozartoz.truffle.nodes.call;

import org.mozartoz.truffle.nodes.OzNode;
import org.mozartoz.truffle.runtime.OzBacktrace;
import org.mozartoz.truffle.runtime.TailCallException;

import com.oracle.truffle.api.CompilerDirectives;
import com.oracle.truffle.api.Truffle;
import com.oracle.truffle.api.frame.FrameDescriptor;
import com.oracle.truffle.api.frame.FrameSlot;
import com.oracle.truffle.api.frame.FrameSlotKind;
import com.oracle.truffle.api.frame.FrameUtil;
import com.oracle.truffle.api.frame.VirtualFrame;
import com.oracle.truffle.api.nodes.LoopNode;
import com.oracle.truffle.api.nodes.RepeatingNode;
import com.oracle.truffle.api.profiles.BranchProfile;
import com.oracle.truffle.api.profiles.LoopConditionProfile;

public class TailCallCatcherNode extends CallableNode {

	@Child CallNode callNode;
	@Child LoopNode loopNode;

	private final BranchProfile normalCallProfile = BranchProfile.create();
	private final BranchProfile tailCallProfile = BranchProfile.create();

	public TailCallCatcherNode(CallNode callNode) {
		this.callNode = callNode;
	}

	public Object execute(VirtualFrame frame) {
		try {
			callNode.execute(frame);
			normalCallProfile.enter();
			return unit;
		} catch (TailCallException tailCall) {
			tailCallProfile.enter();
			return tailCallLoop(frame, tailCall);
		}
	}

	@Override
	public Object executeCall(VirtualFrame frame, Object receiver, Object[] arguments) {
		try {
			callNode.executeCall(frame, receiver, arguments);
			normalCallProfile.enter();
			return unit;
		} catch (TailCallException tailCall) {
			tailCallProfile.enter();
			receiver = null;
			arguments = null;
			return tailCallLoop(frame, tailCall);
		}
	}

	private Object tailCallLoop(VirtualFrame frame, TailCallException tailCall) {
		if (loopNode == null) {
			CompilerDirectives.transferToInterpreterAndInvalidate();
			this.loopNode = insert(Truffle.getRuntime().createLoopNode(new TailCallLoopNode(frame.getFrameDescriptor())));
		}

		TailCallLoopNode tailCallLoop = (TailCallLoopNode) loopNode.getRepeatingNode();
		tailCallLoop.setTailCallException(frame, tailCall);
		loopNode.executeLoop(frame);
		return tailCallLoop.getResult(frame);
	}

	public static class TailCallLoopNode extends OzNode implements RepeatingNode {

		private final FrameSlot tailCallSlot;
		private final LoopConditionProfile loopProfile = LoopConditionProfile.createCountingProfile();
		private final BranchProfile normalCallProfile = BranchProfile.create();
		private final BranchProfile exceptionProfile = BranchProfile.create();

		@Child CallNode callNode = CallNodeGen.create(null, null);

		public TailCallLoopNode(FrameDescriptor frameDescriptor) {
			this.tailCallSlot = frameDescriptor.findOrAddFrameSlot("<OSRtailCall>", FrameSlotKind.Object);
		}

		public void setTailCallException(VirtualFrame frame, TailCallException tailCall) {
			frame.setObject(tailCallSlot, tailCall);
		}

		public Object getResult(VirtualFrame frame) {
			setTailCallException(frame, null);
			return unit;
		}

		@Override
		public boolean executeRepeating(VirtualFrame frame) {
			return loopProfile.profile(loopBody(frame));
		}

		private boolean loopBody(VirtualFrame frame) {
			TailCallException tailCall = (TailCallException) FrameUtil.getObjectSafe(frame, tailCallSlot);
			try {
				callNode.executeCall(frame, tailCall.receiver, tailCall.arguments);
				normalCallProfile.enter();
				return false;
			} catch (TailCallException exception) {
				exceptionProfile.enter();
				setTailCallException(frame, exception);
				return true;
			}
		}

		@Override
		public Object execute(VirtualFrame frame) {
			throw new UnsupportedOperationException();
		}

		@Override
		public String toString() {
			return OzBacktrace.formatNode(this);
		}

	}

}

@@ -2280,24 +2280,65 @@
 ameSlot 
-tailCall
+receiverSlot;%0A%09%09private final FrameSlot arguments
 Slot;%0A%09%09
@@ -2711,24 +2711,131 @@
 %09%09%09this.
-tailCall
+receiverSlot = frameDescriptor.findOrAddFrameSlot(%22%3COSRtailCallReceiver%3E%22, FrameSlotKind.Object);%0A%09%09%09this.arguments
 Slot = f
@@ -2881,16 +2881,20 @@
 tailCall
+Args
 %3E%22, Fram
@@ -2912,16 +2912,47 @@
 bject);%0A
+%09%09%09loopProfile.profile(false);%0A
 %09%09%7D%0A%0A%09%09p
@@ -3052,24 +3052,78 @@
 tObject(
-tailCall
+receiverSlot, tailCall.receiver);%0A%09%09%09frame.setObject(arguments
 Slot, ta
@@ -3124,24 +3124,34 @@
 ot, tailCall
+.arguments
 );%0A%09%09%7D%0A%0A%09%09pu
@@ -3201,34 +3201,77 @@
 %0A%09%09%09
-setTailCallException(frame
+frame.setObject(receiverSlot, null);%0A%09%09%09frame.setObject(argumentsSlot
 , nu
@@ -3473,55 +3473,105 @@
 %0A%09%09%09
-TailCallException tailCall = (TailCallException
+Object receiver = FrameUtil.getObjectSafe(frame, receiverSlot);%0A%09%09%09Object%5B%5D arguments = (Object%5B%5D
 ) Fr
@@ -3599,24 +3599,25 @@
 (frame, 
-tailCall
+arguments
 Slot);%0A%09
@@ -3656,25 +3656,16 @@
 (frame, 
-tailCall.
 receiver
@@ -3666,25 +3666,16 @@
 ceiver, 
-tailCall.
 argument

Save receiver and arguments in 2 different slots to help Escape Analysis