// See the file "COPYING" in the main distribution directory for copyright.

#include "zeek/script_opt/GenIDDefs.h"

#include "zeek/Desc.h"
#include "zeek/Expr.h"
#include "zeek/Reporter.h"
#include "zeek/Scope.h"
#include "zeek/script_opt/ExprOptInfo.h"
#include "zeek/script_opt/ScriptOpt.h"
#include "zeek/script_opt/StmtOptInfo.h"

namespace zeek::detail
	{

GenIDDefs::GenIDDefs(std::shared_ptr<ProfileFunc> _pf, const Func* f, ScopePtr scope, StmtPtr body)
	: pf(std::move(_pf))
	{
	TraverseFunction(f, scope, body);
	}

void GenIDDefs::TraverseFunction(const Func* f, ScopePtr scope, StmtPtr body)
	{
	func_flavor = f->Flavor();

	// Establish the outermost barrior and associated set of
	// identifiers.
	barrier_blocks.push_back(0);
	modified_IDs.push_back({});

	for ( const auto& g : pf->Globals() )
		{
		g->GetOptInfo()->Clear();
		TrackID(g);
		}

	// Clear the locals before processing the arguments, since
	// they're included among the locals.
	for ( const auto& l : pf->Locals() )
		l->GetOptInfo()->Clear();

	const auto& args = scope->OrderedVars();
	int nparam = f->GetType()->Params()->NumFields();

	for ( const auto& a : args )
		{
		if ( --nparam < 0 )
			break;

		a->GetOptInfo()->Clear();
		TrackID(a);
		}

	stmt_num = 0; // 0 = "before the first statement"

	body->Traverse(this);
	}

TraversalCode GenIDDefs::PreStmt(const Stmt* s)
	{
	curr_stmt = s;

	auto si = s->GetOptInfo();
	si->stmt_num = ++stmt_num;
	si->block_level = confluence_blocks.size() + 1;

	switch ( s->Tag() )
		{
		case STMT_CATCH_RETURN:
			{
			auto cr = s->AsCatchReturnStmt();
			auto block = cr->Block();

			StartConfluenceBlock(s);
			block->Traverse(this);
			EndConfluenceBlock();

			auto retvar = cr->RetVar();
			if ( retvar )
				TrackID(retvar->Id());

			return TC_ABORTSTMT;
			}

		case STMT_IF:
			{
			auto i = s->AsIfStmt();
			auto cond = i->StmtExpr();
			auto t_branch = i->TrueBranch();
			auto f_branch = i->FalseBranch();

			cond->Traverse(this);

			StartConfluenceBlock(s);

			t_branch->Traverse(this);
			if ( ! t_branch->NoFlowAfter(false) )
				BranchBeyond(curr_stmt, s, true);

			f_branch->Traverse(this);
			if ( ! f_branch->NoFlowAfter(false) )
				BranchBeyond(curr_stmt, s, true);

			EndConfluenceBlock(true);

			return TC_ABORTSTMT;
			}

		case STMT_SWITCH:
			{
			auto sw = s->AsSwitchStmt();
			auto e = sw->StmtExpr();

			e->Traverse(this);

			StartConfluenceBlock(sw);

			for ( const auto& c : *sw->Cases() )
				{
				auto body = c->Body();

				auto exprs = c->ExprCases();
				if ( exprs )
					exprs->Traverse(this);

				auto type_ids = c->TypeCases();
				if ( type_ids )
					{
					for ( const auto& id : *type_ids )
						if ( id->Name() )
							TrackID(id);
					}

				body->Traverse(this);
				}

			EndConfluenceBlock(sw->HasDefault());

			return TC_ABORTSTMT;
			}

		case STMT_FOR:
			{
			auto f = s->AsForStmt();

			auto ids = f->LoopVars();
			auto e = f->LoopExpr();
			auto body = f->LoopBody();
			auto val_var = f->ValueVar();

			e->Traverse(this);

			for ( const auto& id : *ids )
				TrackID(id);

			if ( val_var )
				TrackID(val_var);

			StartConfluenceBlock(s);
			body->Traverse(this);

			if ( ! body->NoFlowAfter(false) )
				BranchBackTo(curr_stmt, s, true);

			EndConfluenceBlock();

			return TC_ABORTSTMT;
			}

		case STMT_WHILE:
			{
			auto w = s->AsWhileStmt();

			StartConfluenceBlock(s);

			auto cond_pred_stmt = w->CondPredStmt();
			if ( cond_pred_stmt )
				cond_pred_stmt->Traverse(this);

			// Important to traverse the condition in its version
			// interpreted as a statement, so that when evaluating
			// its variable usage, that's done in the context of
			// *after* cond_pred_stmt executes, rather than as
			// part of that execution.
			auto cond_stmt = w->ConditionAsStmt();
			cond_stmt->Traverse(this);

			auto body = w->Body();
			body->Traverse(this);

			if ( ! body->NoFlowAfter(false) )
				BranchBackTo(curr_stmt, s, true);

			EndConfluenceBlock();

			return TC_ABORTSTMT;
			}

		case STMT_WHEN:
			{
			// ### punt on these for now, need to reflect on bindings.
			return TC_ABORTSTMT;
			}

		default:
			return TC_CONTINUE;
		}
	}

TraversalCode GenIDDefs::PostStmt(const Stmt* s)
	{
	switch ( s->Tag() )
		{
		case STMT_INIT:
			{
			auto init = s->AsInitStmt();
			auto& inits = init->Inits();

			for ( const auto& id : inits )
				{
				auto id_t = id->GetType();

				// Only aggregates get initialized.
				if ( zeek::IsAggr(id->GetType()->Tag()) )
					TrackID(id);
				}

			break;
			}

		case STMT_RETURN:
			ReturnAt(s);
			break;

		case STMT_NEXT:
			BranchBackTo(curr_stmt, FindLoop(), false);
			break;

		case STMT_BREAK:
			{
			auto target = FindBreakTarget();

			if ( target )
				BranchBeyond(s, target, false);

			else
				{
				ASSERT(func_flavor == FUNC_FLAVOR_HOOK);
				ReturnAt(s);
				}

			break;
			}

		case STMT_FALLTHROUGH:
			// No need to do anything, the work all occurs
			// with NoFlowAfter.
			break;

		default:
			break;
		}

	return TC_CONTINUE;
	}

TraversalCode GenIDDefs::PreExpr(const Expr* e)
	{
	e->GetOptInfo()->stmt_num = stmt_num;

	switch ( e->Tag() )
		{
		case EXPR_NAME:
			CheckVarUsage(e, e->AsNameExpr()->Id());
			break;

		case EXPR_ASSIGN:
			{
			auto lhs = e->GetOp1();
			auto op2 = e->GetOp2();

			if ( lhs->Tag() == EXPR_LIST && op2->GetType()->Tag() != TYPE_ANY )
				{
				// This combination occurs only for assignments used
				// to initialize table entries.  Treat it as references
				// to both the lhs and the rhs, not as an assignment.
				return TC_CONTINUE;
				}

			op2->Traverse(this);

			if ( ! CheckLHS(lhs, op2) )
				// Not a simple assignment (or group of assignments),
				// so analyze the accesses to check for use of
				// possibly undefined values.
				lhs->Traverse(this);

			return TC_ABORTSTMT;
			}

		case EXPR_COND:
			// Special hack.  We turn off checking for usage issues
			// inside conditionals.  This is because we use them heavily
			// to deconstruct logical expressions for which the actual
			// operand access is safe (guaranteed not to access a value
			// that hasn't been undefined), but the flow analysis has
			// trouble determining that.
			++suppress_usage;
			e->GetOp1()->Traverse(this);
			e->GetOp2()->Traverse(this);
			e->GetOp3()->Traverse(this);
			--suppress_usage;

			return TC_ABORTSTMT;

		case EXPR_LAMBDA:
			{
			auto l = static_cast<const LambdaExpr*>(e);
			const auto& ids = l->OuterIDs();

			for ( auto& id : ids )
				CheckVarUsage(e, id);

			// Don't descend into the lambda body - we'll analyze and
			// optimize it separately, as its own function.
			return TC_ABORTSTMT;
			}

		default:
			break;
		}

	return TC_CONTINUE;
	}

TraversalCode GenIDDefs::PostExpr(const Expr* e)
	{
	// Attend to expressions that reflect assignments after
	// execution, but for which the assignment target was
	// also an accessed value (so if we analyzed them
	// in PreExpr then we'd have had to do manual traversals
	// of their operands).

	auto t = e->Tag();
	if ( t == EXPR_INCR || t == EXPR_DECR || t == EXPR_ADD_TO || t == EXPR_REMOVE_FROM )
		{
		auto op = e->GetOp1();
		if ( ! IsAggr(op) )
			(void)CheckLHS(op);
		}

	return TC_CONTINUE;
	}

bool GenIDDefs::CheckLHS(const ExprPtr& lhs, const ExprPtr& rhs)
	{
	switch ( lhs->Tag() )
		{
		case EXPR_REF:
			return CheckLHS(lhs->GetOp1(), rhs);

		case EXPR_NAME:
			{
			auto n = lhs->AsNameExpr();
			TrackID(n->Id(), rhs);
			return true;
			}

		case EXPR_LIST:
			{ // look for [a, b, c] = any_val
			auto l = lhs->AsListExpr();
			for ( const auto& expr : l->Exprs() )
				{
				if ( expr->Tag() != EXPR_NAME )
					// This will happen for table initializers,
					// for example.
					return false;

				auto n = expr->AsNameExpr();
				TrackID(n->Id());
				}

			return true;
			}

		case EXPR_FIELD:
			// If we want to track record field initializations,
			// we'd handle that here.
			return false;

		case EXPR_INDEX:
			// If we wanted to track potential alterations of
			// aggregates, we'd do that here.
			return false;

		default:
			reporter->InternalError("bad tag in GenIDDefs::CheckLHS");
		}
	}

bool GenIDDefs::IsAggr(const Expr* e) const
	{
	if ( e->Tag() != EXPR_NAME )
		return false;

	auto n = e->AsNameExpr();
	auto id = n->Id();
	auto tag = id->GetType()->Tag();

	return zeek::IsAggr(tag);
	}

void GenIDDefs::CheckVarUsage(const Expr* e, const ID* id)
	{
	if ( analysis_options.usage_issues != 1 || id->IsGlobal() || suppress_usage > 0 )
		return;

	auto oi = id->GetOptInfo();

	if ( ! oi->DidUndefinedWarning() && ! oi->IsDefinedBefore(curr_stmt) &&
	     ! id->GetAttr(ATTR_IS_ASSIGNED) )
		{
		if ( ! oi->IsPossiblyDefinedBefore(curr_stmt) )
			{
			e->Warn("used without definition");
			oi->SetDidUndefinedWarning();
			}

		else if ( ! oi->DidPossiblyUndefinedWarning() )
			{
			e->Warn("possibly used without definition");
			oi->SetDidPossiblyUndefinedWarning();
			}
		}
	}

void GenIDDefs::StartConfluenceBlock(const Stmt* s)
	{
	if ( s->Tag() == STMT_CATCH_RETURN )
		barrier_blocks.push_back(confluence_blocks.size());

	confluence_blocks.push_back(s);
	modified_IDs.push_back({});
	}

void GenIDDefs::EndConfluenceBlock(bool no_orig)
	{
	for ( auto id : modified_IDs.back() )
		id->GetOptInfo()->ConfluenceBlockEndsAfter(curr_stmt, no_orig);

	confluence_blocks.pop_back();

	auto bb = barrier_blocks.back();
	if ( bb > 0 && confluence_blocks.size() == bb )
		barrier_blocks.pop_back();

	modified_IDs.pop_back();
	}

void GenIDDefs::BranchBackTo(const Stmt* from, const Stmt* to, bool close_all)
	{
	for ( auto id : modified_IDs.back() )
		id->GetOptInfo()->BranchBackTo(from, to, close_all);
	}

void GenIDDefs::BranchBeyond(const Stmt* from, const Stmt* to, bool close_all)
	{
	for ( auto id : modified_IDs.back() )
		id->GetOptInfo()->BranchBeyond(from, to, close_all);

	to->GetOptInfo()->contains_branch_beyond = true;
	}

const Stmt* GenIDDefs::FindLoop()
	{
	int i = confluence_blocks.size() - 1;
	while ( i >= 0 )
		{
		auto t = confluence_blocks[i]->Tag();
		if ( t == STMT_WHILE || t == STMT_FOR )
			break;

		--i;
		}

	ASSERT(i >= 0);

	return confluence_blocks[i];
	}

const Stmt* GenIDDefs::FindBreakTarget()
	{
	int i = confluence_blocks.size() - 1;
	while ( i >= 0 )
		{
		auto cb = confluence_blocks[i];
		auto t = cb->Tag();
		if ( t == STMT_WHILE || t == STMT_FOR || t == STMT_SWITCH )
			return cb;

		--i;
		}

	return nullptr;
	}

void GenIDDefs::ReturnAt(const Stmt* s)
	{
	for ( auto id : modified_IDs.back() )
		id->GetOptInfo()->ReturnAt(s);
	}

void GenIDDefs::TrackID(const ID* id, const ExprPtr& e)
	{
	auto oi = id->GetOptInfo();

	ASSERT(! barrier_blocks.empty());
	oi->DefinedAfter(curr_stmt, e, confluence_blocks, barrier_blocks.back());

	// Ensure we track this identifier across all relevant
	// confluence regions.
	for ( auto i = barrier_blocks.back(); i < confluence_blocks.size(); ++i )
		// Add one because modified_IDs includes outer non-confluence
		// block.
		modified_IDs[i + 1].insert(id);

	if ( confluence_blocks.empty() )
		// This is a definition at the outermost level.
		modified_IDs[0].insert(id);
	}

	} // zeek::detail
