//
// Cforall Version 1.0.0 Copyright (C) 2015 University of Waterloo
//
// The contents of this file are covered under the licence agreement in the
// file "LICENCE" distributed with Cforall.
//
// Corun.cpp -- generate code needed by the actor system
//
// Author           : Colby Parsons
// Created On       : Monday October 9 15:16:42 2023
// Last Modified By : Peter A. Buhr
// Last Modified On : Thu Dec 14 17:32:17 2023
// Update Count     : 1
//

#include "AST/Decl.hpp"
#include "AST/Expr.hpp"
#include "AST/Pass.hpp"
#include "AST/Stmt.hpp"
#include "AST/TranslationUnit.hpp"
#include "Common/UniqueName.hpp"
using namespace ast;
using namespace std;

namespace Concurrency {

struct CorunKeyword : public WithDeclsToAdd, public WithStmtsToAdd {
	UniqueName CorunFnNamer = "__CFA_corun_lambda_"s;
	UniqueName CoforFnNamer = "__CFA_cofor_lambda_"s;
	// UniqueName CoforFnVarNamer = "__CFA_cofor_lambda_var"s;
	UniqueName RunnerBlockNamer = "__CFA_corun_block_"s;

	string coforArgName = "__CFA_cofor_lambda_arg";
	string numProcsName = "__CFA_cofor_num_procs";
	string currProcsName = "__CFA_cofor_curr_procs";
	string thdArrName = "__CFA_cofor_thread_array";
	string loopTempName = "__CFA_cofor_loop_temp";


	const StructDecl * runnerBlockDecl = nullptr;
	const StructDecl * coforRunnerDecl = nullptr;

	// Finds runner_block (corun task) and cofor_runner (cofor task) decls
	void previsit( const StructDecl * decl ) {
		if ( !decl->body ) {
			return;
		} else if ( "runner_block" == decl->name ) {
			assert( !runnerBlockDecl );
			runnerBlockDecl = decl;
		} else if ( "cofor_runner" == decl->name ) {
			assert( !coforRunnerDecl );
			coforRunnerDecl = decl;
		}
	}

	// codegen for cofor statements
	Stmt * postvisit( const CoforStmt * stmt ) {
		if ( !runnerBlockDecl || !coforRunnerDecl )
			SemanticError( stmt->location, "To use cofor statements add #include <cofor.hfa>" );

		if ( stmt->inits.size() != 1 )
			SemanticError( stmt->location, "Cofor statements must have a single initializer in the loop control" );

		if ( !stmt->body )
			return nullptr;

		const CodeLocation & loc = stmt->location;
		const string fnName = CoforFnNamer.newName();

		CompoundStmt * body = new CompoundStmt( loc );

		// push back cofor initializer to generated body
		body->push_back( deepCopy( stmt->inits.at(0) ) );

		CompoundStmt * fnBody = new CompoundStmt( loc );

		const DeclStmt * declStmtPtr = dynamic_cast<const DeclStmt *>(stmt->inits.at(0).get());
		if ( ! declStmtPtr )
			SemanticError( stmt->location, "Cofor statement initializer is somehow not a decl statement?" );

		const Decl * declPtr = dynamic_cast<const Decl *>(declStmtPtr->decl.get());
		if ( ! declPtr )
			SemanticError( stmt->location, "Cofor statement initializer is somehow not a decl?" );

		Type * initType = new TypeofType( new NameExpr( loc, declPtr->name ) );

		// Generates:
		// typeof(init) __CFA_cofor_lambda_var = *((typeof(init) *)val);
		fnBody->push_back( new DeclStmt( loc,
			new ObjectDecl( loc,
				declPtr->name,
				initType,
				new SingleInit( loc,
					UntypedExpr::createDeref( loc,
						new CastExpr( loc,
							new NameExpr( loc, coforArgName ),
							new PointerType( initType ), ExplicitCast
						)
					)
				)
			)
		));

		// push rest of cofor body into loop lambda
		fnBody->push_back( deepCopy( stmt->body ) );

		// Generates:
		// void __CFA_cofor_lambda_() {
		//    typeof(init) __CFA_cofor_lambda_var = *((typeof(init) *)val);
		//    stmt->body;
		// }
		Stmt * coforLambda = new DeclStmt( loc,
			new FunctionDecl( loc,
				fnName,                                             // name
				{
					new ObjectDecl( loc,
						coforArgName,
						new ast::PointerType( new ast::VoidType() )
					)
				},                                                  // params
				{},                                                 // return
				fnBody   // body
			)
		);
		body->push_back( coforLambda );

		// Generates:
		// unsigned __CFA_cofor_num_procs = get_proc_count();
		body->push_back( new DeclStmt( loc,
				new ObjectDecl( loc,
					numProcsName,
					new BasicType( BasicKind::UnsignedInt ),
					new SingleInit( loc,
						new UntypedExpr( loc,
							new NameExpr( loc, "get_proc_count" ),
							{}
						)
					)
				)
			)
		);

		// Generates:
		// unsigned __CFA_cofor_curr_procs = 0;
		body->push_back( new DeclStmt( loc,
				new ObjectDecl( loc,
					currProcsName,
					new BasicType( BasicKind::UnsignedInt ),
					new SingleInit( loc, ConstantExpr::from_int( loc, 0 ) )
				)
			)
		);

		// Generates:
		// unsigned cofor_runner __CFA_cofor_thread_array[nprocs];
		body->push_back( new DeclStmt( loc,
				new ObjectDecl( loc,
					thdArrName,
					new ast::ArrayType(
						new StructInstType( coforRunnerDecl ),
						new NameExpr( loc, numProcsName ),
						ast::FixedLen,
						ast::DynamicDim
					)
				)
			)
		);

		// Generates:
		// start_runners( __CFA_cofor_thread_array, __CFA_cofor_num_procs, __CFA_cofor_lambda_ );
		body->push_back( new ExprStmt( loc,
			new UntypedExpr( loc,
				new NameExpr( loc, "start_runners" ),
				{
					new NameExpr( loc, thdArrName ),
					new NameExpr( loc, numProcsName ),
					new NameExpr( loc, fnName )
				}
			)
		));

		// Generates:
		// typeof(initializer) * __CFA_cofor_loop_temp = malloc();
		CompoundStmt * forLoopBody = new CompoundStmt( loc );
		forLoopBody->push_back( new DeclStmt( loc,
				new ObjectDecl( loc,
					loopTempName,
					new PointerType( initType ),
					new SingleInit( loc,
						new UntypedExpr( loc,
							new NameExpr( loc, "malloc" ),
							{}
						)
					)
				)
			)
		);

		// Generates:
		// *__CFA_cofor_loop_temp = initializer;
		forLoopBody->push_back( new ExprStmt( loc,
			UntypedExpr::createAssign( loc,
				UntypedExpr::createDeref( loc, new NameExpr( loc, loopTempName ) ),
				new NameExpr( loc, declPtr->name )
			)
		));

		// Generates:
		// send_work( __CFA_cofor_thread_array, __CFA_cofor_num_procs,
		//     __CFA_cofor_curr_procs, __CFA_cofor_loop_temp );
		forLoopBody->push_back( new ExprStmt( loc,
			new UntypedExpr( loc,
				new NameExpr( loc, "send_work" ),
				{
					new NameExpr( loc, thdArrName ),
					new NameExpr( loc, numProcsName ),
					new NameExpr( loc, currProcsName ),
					new NameExpr( loc, loopTempName )
				}
			)
		));

		body->push_back( new ForStmt( loc,
			{},
			deepCopy( stmt->cond ),
			deepCopy( stmt->inc ),
			forLoopBody
		));

		// Generates:
		// end_runners( __CFA_cofor_thread_array, __CFA_cofor_num_procs );
		body->push_back( new ExprStmt( loc,
			new UntypedExpr( loc,
				new NameExpr( loc, "end_runners" ),
				{
					new NameExpr( loc, thdArrName ),
					new NameExpr( loc, numProcsName )
				}
			)
		));

		return body;
	}

	// codegen for corun statements
	Stmt * postvisit( const CorunStmt * stmt ) {
		if ( !runnerBlockDecl || !coforRunnerDecl )
			SemanticError( stmt->location, "To use corun statements add #include <cofor.hfa>" );

		if ( !stmt->stmt )
			return nullptr;

		const CodeLocation & loc = stmt->location;
		const string fnName = CorunFnNamer.newName();
		const string objName = RunnerBlockNamer.newName();

		// Generates:
		// void __CFA_corun_lambda_() { ... stmt->stmt ... }
		Stmt * runnerLambda = new DeclStmt( loc,
			new FunctionDecl( loc,
				fnName,                                             // name
				{},                                                 // params
				{},                                                 // return
				new CompoundStmt( loc, { deepCopy(stmt->stmt) } )   // body
			)
		);

		// Generates:
		// runner_block __CFA_corun_block_;
		Stmt * objDecl = new DeclStmt( loc,
			new ObjectDecl( loc,
				objName,
				new StructInstType( runnerBlockDecl )
			)
		);

		// Generates:
		// __CFA_corun_block_{ __CFA_corun_lambda_ };
		Stmt * threadStart = new ExprStmt( loc,
			new UntypedExpr ( loc,
				new NameExpr( loc, "?{}" ),
				{
					new NameExpr( loc, objName ),
					new NameExpr( loc, fnName )
				}
			)
		);

		stmtsToAddBefore.push_back( runnerLambda );
		stmtsToAddBefore.push_back( objDecl );

		return threadStart;
	}
};

void implementCorun( TranslationUnit & translationUnit ) {
	Pass<CorunKeyword>::run( translationUnit );
}

} // namespace Concurrency