%-----------------------------------------------------------------------------%
% Copyright (C) 1999 The University of Melbourne.
% This file may only be copied under the terms of the GNU General
% Public License - see the file COPYING in the Mercury distribution.
%-----------------------------------------------------------------------------%

% File: ml_tailcall.m
% Main author: fjh

% This module is an MLDS-to-MLDS transformation
% that marks function calls as tail calls whenever
% it is safe to do so, based on the assumptions described below.

% A function call can safely be marked as a tail call if
%	(1) it occurs in a position which would fall through into the
%	    end of the function body or to a `return' statement,
% and
%	(2) the lvalues in which the return value(s) from the `call'
%	    will be placed are the same as the value(s) returned
%	    by the `return', and these lvalues are all local variables.
% and
%	(3) the function's local variables do not need to be live
%	    for that call.
%
% For (2), we just assume (rather than checking) that 
% any variables returned by the `return' statement are
% local variables.  This assumption is true for the MLDS code
% generated by ml_code_gen.m.
%
% For (3), we assume that the addresses of local variables
% and nested functions are only ever passed down to other functions
% (and used to assign to the local variable or to call the nested
% function), so that here we only need to check if the potential
% tail call uses such addresses, not whether such addresses were
% taken in earlier calls.  That is, if the addresses
% of locals were taken in earlier calls from the same function,
% we assume that these addresses will not be saved (on the heap,
% or in global variables, etc.) and used after those earlier calls
% have returned.  This assumption is true for the MLDS code generated
% by ml_code_gen.m.

%-----------------------------------------------------------------------------%

:- module ml_tailcall.
:- interface.
:- import_module mlds.
:- import_module io.

	% Traverse the MLDS, marking all optimizable tail calls
	% as tail calls.
	%
:- pred ml_mark_tailcalls(mlds, mlds, io__state, io__state).
:- mode ml_mark_tailcalls(in, out, di, uo) is det.

%-----------------------------------------------------------------------------%
%-----------------------------------------------------------------------------%

:- implementation.
:- import_module list, std_util.

ml_mark_tailcalls(MLDS0, MLDS) -->
	{ MLDS0 = mlds(ModuleName, ForeignCode, Imports, Defns0) },
	{ Defns = mark_tailcalls_in_defns(Defns0) },
	{ MLDS = mlds(ModuleName, ForeignCode, Imports, Defns) }.

%-----------------------------------------------------------------------------%

	% The `at_tail' type indicates whether or not a subgoal
	% is at a tail position, i.e. is followed by a return
	% statement or the end of the function, and if so,
	% specifies the return values (if any) in the return statement.
:- type at_tail == maybe(list(mlds__rval)).

	% The `locals' type contains a list of local definitions
	% which are in scope.
:- type locals == list(local_defns).
:- type local_defns
	--->	params(mlds__arguments)		
	;	defns(mlds__defns)
	.

%-----------------------------------------------------------------------------%

%
% mark_tailcalls_in_defns:
% mark_tailcalls_in_defn:
%	Recursively process the definition(s),
%	marking each optimizable tail call in them as a tail call.
%
% mark_tailcalls_in_maybe_statement:
% mark_tailcalls_in_statements:
% mark_tailcalls_in_statement:
% mark_tailcalls_in_stmt:
% mark_tailcalls_in_atomic_stmt:
%	Recursively process the statement(s),
%	marking each optimizable tail call in them as a tail call.
%	The `AtTail' argument indicates whether or not this
%	construct is in a tail call position.
%	The `Locals' argument contains a list of the
%	local definitions which are in scope at this point.
%

:- func mark_tailcalls_in_defns(mlds__defns) = mlds__defns.
mark_tailcalls_in_defns(Defns) = list__map(mark_tailcalls_in_defn, Defns).

:- func mark_tailcalls_in_defn(mlds__defn) = mlds__defn.
mark_tailcalls_in_defn(Defn0) = Defn :-
	Defn0 = mlds__defn(Name, Context, Flags, DefnBody0),
	(
		DefnBody0 = mlds__function(PredProcId, Params, FuncBody0),
		%
		% Compute the initial value of the `Locals' and
		% `AtTail' arguments.
		%
		Params = mlds__func_params(Args, RetTypes),
		Locals = [params(Args)],
		( RetTypes = [] ->
			AtTail = yes([])
		;
			AtTail = no
		),
		FuncBody = mark_tailcalls_in_maybe_statement(FuncBody0,
			AtTail, Locals),
		DefnBody = mlds__function(PredProcId, Params, FuncBody),
		Defn = mlds__defn(Name, Context, Flags, DefnBody)
	;
		DefnBody0 = mlds__data(_, _),
		Defn = Defn0
	;
		DefnBody0 = mlds__class(ClassDefn0),
		ClassDefn0 = class_defn(Kind, Imports, BaseClasses, Implements,
		                MemberDefns0),
		MemberDefns = mark_tailcalls_in_defns(MemberDefns0),
		ClassDefn = class_defn(Kind, Imports, BaseClasses, Implements,
		                MemberDefns),
		DefnBody = mlds__class(ClassDefn),
		Defn = mlds__defn(Name, Context, Flags, DefnBody)
	).

:- func mark_tailcalls_in_maybe_statement(maybe(mlds__statement),
		at_tail, locals) = maybe(mlds__statement).

mark_tailcalls_in_maybe_statement(no, _, _) = no.
mark_tailcalls_in_maybe_statement(yes(Statement0), AtTail, Locals) =
		yes(Statement) :-
	Statement = mark_tailcalls_in_statement(Statement0, AtTail, Locals).
	
:- func mark_tailcalls_in_statements(mlds__statements, at_tail, locals) =
		mlds__statements.

mark_tailcalls_in_statements([], _, _) = [].
mark_tailcalls_in_statements([First0 | Rest0], AtTail, Locals) =
		[First | Rest] :-
	%
	% If the First statement is followed by a `return'
	% statement, then it is in a tailcall position.
	% If there are no statements after the first, then
	% the first statement is in a tail call position
	% iff the statement list is in a tail call position.
	% Otherwise, i.e. if the first statement is followed
	% by anything other than a `return' statement, then
	% the first statement is not in a tail call position.
	%
	(
		Rest = [mlds__statement(return(ReturnVals), _) | _]
	->
		FirstAtTail = yes(ReturnVals)
	;
		Rest = []
	->
		FirstAtTail = AtTail
	;
		FirstAtTail = no
	),
	First = mark_tailcalls_in_statement(First0, FirstAtTail, Locals),
	Rest = mark_tailcalls_in_statements(Rest0, AtTail, Locals).

:- func mark_tailcalls_in_statement(mlds__statement, at_tail, locals) =
		mlds__statement.

mark_tailcalls_in_statement(Statement0, AtTail, Locals) = Statement :-
	Statement0 = mlds__statement(Stmt0, Context),
	Stmt = mark_tailcalls_in_stmt(Stmt0, AtTail, Locals),
	Statement = mlds__statement(Stmt, Context).

:- func mark_tailcalls_in_stmt(mlds__stmt, at_tail, locals) = mlds__stmt.

mark_tailcalls_in_stmt(Stmt0, AtTail, Locals) = Stmt :-
	(
		%
		% Whenever we encounter a block statement,
		% we recursively mark tailcalls in any nested
		% functions defined in that block.
		% We also need to add any local definitions in that
		% block to the list of currently visible local
		% declarations before processing the statements
		% in that block.  The statement list will be in a
		% tail position iff the block is in a tail position.
		%
		Stmt0 = block(Defns0, Statements0),
		Defns = mark_tailcalls_in_defns(Defns0),
		NewLocals = [defns(Defns) | Locals],
		Statements = mark_tailcalls_in_statements(Statements0,
				AtTail, NewLocals),
		Stmt = block(Defns, Statements)
	;
		%
		% The statement in the body of a while loop is never
		% in a tail position.
		%
		Stmt0 = while(Rval, Statement0, Once),
		Statement = mark_tailcalls_in_statement(Statement0, no, Locals),
		Stmt = while(Rval, Statement, Once)
	;
		%
		% Both the `then' and the `else' parts of an if-then-else
		% are in a tail position iff the if-then-else is in a
		% tail position.
		%
		Stmt0 = if_then_else(Cond, Then0, MaybeElse0),
		Then = mark_tailcalls_in_statement(Then0, AtTail, Locals),
		MaybeElse = mark_tailcalls_in_maybe_statement(MaybeElse0,
				AtTail, Locals),
		Stmt = if_then_else(Cond, Then, MaybeElse)
	;
		Stmt0 = label(_),
		Stmt = Stmt0
	;
		Stmt0 = goto(_),
		Stmt = Stmt0
	;
		Stmt0 = computed_goto(_, _),
		Stmt = Stmt0
	;
		Stmt0 = call(Sig, Func, Obj, Args, ReturnLvals, _TailCall0),
		%
		% check if we can mark this call as a tail call
		%
		(
			%
			% we must be in a tail position
			%
			AtTail = yes(ReturnRvals),
			%
			% the values returned in this call must match
			% those returned by the `return' statement that
			% follows
			%
			match_return_vals(ReturnRvals, ReturnLvals),
			%
			% the call must not take the address of any
			% local variables or nested functions
			%
			check_maybe_rval(Obj, Locals),
			check_rvals(Args, Locals)
		->
			% mark this call as a tail call
			TailCall = tail_call,
			Stmt = call(Sig, Func, Obj, Args, ReturnLvals,
					TailCall)
		;
			% leave this call unchanged
			Stmt = Stmt0
		)
	;
		Stmt0 = return(_Rvals),
		Stmt = Stmt0
	;
		Stmt0 = do_commit(_Ref),
		Stmt = Stmt0
	;
		Stmt0 = try_commit(Ref, Statement0, Handler0),
		%
		% Both the statement inside a `try_commit' and the
		% handler are in tail call position iff the
		% `try_commit' statement is in a tail call position.
		%
		Statement = mark_tailcalls_in_statement(Statement0, AtTail,
				Locals),
		Handler = mark_tailcalls_in_statement(Handler0, AtTail, Locals),
		Stmt = try_commit(Ref, Statement, Handler)
	;
		Stmt0 = atomic(_),
		Stmt = Stmt0
	).

%-----------------------------------------------------------------------------%

%
% match_return_vals(Rvals, Lvals):
% match_return_val(Rval, Lval):
%	Check that the Lval(s) returned by a call match
%	the Rval(s) in the `return' statement that follows,
%	and those Lvals are local variables
%	(so that assignments to them won't have any side effects),
%	so that we can optimize the call into a tailcall.
%

:- pred match_return_vals(list(mlds__rval), list(mlds__lval)).
:- mode match_return_vals(in, in) is semidet.

match_return_vals([], []).
match_return_vals([Rval|Rvals], [Lval|Lvals]) :-
	match_return_val(Rval, Lval),
	match_return_vals(Rvals, Lvals).

:- pred match_return_val(mlds__rval, mlds__lval).
:- mode match_return_val(in, in) is semidet.

match_return_val(lval(Lval), Lval) :-
	lval_is_local(Lval).

:- pred lval_is_local(mlds__lval).
:- mode lval_is_local(in) is semidet.

lval_is_local(var(_)) :-
	% We just assume it is local.  (This assumption is
	% true for the code generated by ml_code_gen.m.)
	true.
lval_is_local(field(_Tag, Rval, _Field)) :-
	% a field of a local variable is local
	( Rval = mem_addr(Lval) ->
		lval_is_local(Lval)
	;
		fail
	).
lval_is_local(mem_ref(_Rval)) :-
	fail.

%-----------------------------------------------------------------------------%

%
% check_rvals:
% check_maybe_rval:
% check_rval:
%	Fail if the specified rval(s) might evaluate to the addresses of
%	local variables (or fields of local variables) or nested functions.
%

:- pred check_rvals(list(mlds__rval), locals).
:- mode check_rvals(in, in) is semidet.

check_rvals([], _).
check_rvals([Rval|Rvals], Locals) :-
	check_rval(Rval, Locals),
	check_rvals(Rvals, Locals).

:- pred check_maybe_rval(maybe(mlds__rval), locals).
:- mode check_maybe_rval(in, in) is semidet.

check_maybe_rval(no, _).
check_maybe_rval(yes(Rval), Locals) :-
	check_rval(Rval, Locals).

:- pred check_rval(mlds__rval, locals).
:- mode check_rval(in, in) is semidet.

check_rval(lval(_Lval), _) :-
	% Passing the _value_ of an lval is fine.
	true.
check_rval(mkword(_Tag, Rval), Locals) :-
	check_rval(Rval, Locals).
check_rval(const(Const), Locals) :-
	check_const(Const, Locals).
check_rval(unop(_Op, Rval), Locals) :-
	check_rval(Rval, Locals).
check_rval(binop(_Op, X, Y), Locals) :-
	check_rval(X, Locals),
	check_rval(Y, Locals).
check_rval(mem_addr(Lval), Locals) :-
	% Passing the address of an lval is a problem,
	% if that lval names a local variable.
	check_lval(Lval, Locals).

%
% check_lval:
%	Fail if the specified lval might be a local variable
%	(or a field of a local variable).
%

:- pred check_lval(mlds__lval, locals).
:- mode check_lval(in, in) is semidet.

check_lval(field(_MaybeTag, Rval, _FieldId), Locals) :-
	check_rval(Rval, Locals).
check_lval(mem_ref(_), _) :-
	% We assume that the addresses of local variables are only
	% ever passed down to other functions, or assigned to,
	% so a mem_ref lval can never refer to a local variable.
	true.
check_lval(var(Var0), Locals) :-
	\+ var_is_local(Var0, Locals).

%
% check_const:
%	Fail if the specified const might be the address of a
%	local variable or nested function.
%
%	The addresses of local variables are probably
%	not consts, at least not unless those variables are
%	declared as static (i.e. `one_copy'),
%	so it might be safe to allow all data_addr_consts here,
%	but currently we just take a conservative approach.
%

:- pred check_const(mlds__rval_const, locals).
:- mode check_const(in, in) is semidet.

check_const(Const, Locals) :-
	( Const = code_addr_const(CodeAddr) ->
		\+ function_is_local(CodeAddr, Locals)
	; Const = data_addr_const(DataAddr) ->
		DataAddr = data_addr(ModuleName, DataName),
		( DataName = var(VarName) ->
			\+ var_is_local(qual(ModuleName, VarName), Locals)
		;
			true
		)
	;
		true
	).

%
% var_is_local:
%	Check whether the specified variable is defined locally,
%	i.e. in storage that might no longer exist when the function
%	returns or does a tail call.
%
%	It would be safe to fail for variables declared static
%	(i.e. `one_copy'), but currently we just take a conservative
% 	approach.
%

:- pred var_is_local(mlds__var, locals).
:- mode var_is_local(in, in) is semidet.

var_is_local(Var, Locals) :-
		% XXX we ignore the ModuleName --
		% that is safe, but overly conservative
	Var = qual(_ModuleName, VarName),
	some [Local] (
		locals_member(Local, Locals),
		Local = data(var(VarName))
	).

%
% function_is_local:
%	Check whether the specified function is defined locally
%	(i.e. as a nested function).
%
:- pred function_is_local(mlds__code_addr, locals).
:- mode function_is_local(in, in) is semidet.

function_is_local(CodeAddr, Locals) :-
	(	
		CodeAddr = proc(QualifiedProcLabel),
	  	MaybeSeqNum = no
	;
		CodeAddr = internal(QualifiedProcLabel, SeqNum),
	  	MaybeSeqNum = yes(SeqNum)
	),
		% XXX we ignore the ModuleName --
		% that is safe, but might be overly conservative
	QualifiedProcLabel = qual(_ModuleName, ProcLabel),
	ProcLabel = PredLabel - ProcId,
	some [Local] (
		locals_member(Local, Locals),
		Local = function(PredLabel, ProcId, MaybeSeqNum, _PredId)
	).

%
% locals_member(Name, Locals):
%	Nondeterministically enumerates the names of all the entities
%	in Locals.
%

:- pred locals_member(mlds__entity_name, locals).
:- mode locals_member(out, in) is nondet.

locals_member(Name, LocalsList) :-
	list__member(Locals, LocalsList),
	(
		Locals = defns(Defns),
		list__member(Defn, Defns),
		Defn = mlds__defn(Name, _, _, _)
	;
		Locals = params(Params),
		list__member(Param, Params),
		Param = Name - _Type
	).

%-----------------------------------------------------------------------------%