//
// 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.
//
// ResolveAssertions.cc --
//
// Author           : Aaron B. Moss
// Created On       : Fri Oct 05 13:46:00 2018
// Last Modified By : Aaron B. Moss
// Last Modified On : Fri Oct 05 13:46:00 2018
// Update Count     : 1
//

#include "ResolveAssertions.h"

#include <algorithm>                // for sort
#include <cassert>                  // for assertf
#include <list>                     // for list
#include <memory>                   // for unique_ptr
#include <unordered_map>            // for unordered_map, unordered_multimap
#include <utility>                  // for move
#include <vector>                   // for vector

#include "Alternative.h"            // for Alternative, AssertionItem, AssertionList
#include "Common/FilterCombos.h"    // for filterCombos
#include "Common/utility.h"         // for sort_mins
#include "ResolvExpr/RenameVars.h"  // for renameTyVars
#include "SymTab/Indexer.h"         // for Indexer
#include "SymTab/Mangler.h"         // for Mangler
#include "SynTree/Expression.h"     // for InferredParams
#include "TypeEnvironment.h"        // for TypeEnvironment, etc.
#include "typeops.h"                // for adjustExprType
#include "Unify.h"                  // for unify

namespace ResolvExpr {
	/// Unified assertion candidate
	struct AssnCandidate {
		SymTab::Indexer::IdData cdata;  ///< Satisfying declaration
		Type* adjType;                  ///< Satisfying type
		TypeEnvironment env;            ///< Post-unification environment
		AssertionSet have;              ///< Post-unification have-set
		AssertionSet need;              ///< Post-unification need-set
		OpenVarSet openVars;            ///< Post-unification open-var set
		UniqueId resnSlot;              ///< Slot for any recursive assertion IDs

		AssnCandidate( const SymTab::Indexer::IdData& cdata, Type* adjType, TypeEnvironment&& env, 
			AssertionSet&& have, AssertionSet&& need, OpenVarSet&& openVars, UniqueId resnSlot ) 
		: cdata(cdata), adjType(adjType), env(std::move(env)), have(std::move(have)), 
			need(std::move(need)), openVars(std::move(openVars)), resnSlot(resnSlot) {}
	};

	/// List of candidate assertion resolutions
	using CandidateList = std::vector<AssnCandidate>;

	/// Reference to single deferred item
	struct DeferRef {
		const DeclarationWithType* decl;
		const AssertionSetValue& info;
		const AssnCandidate& match;
	};

	/// Wrapper for the deferred items from a single assertion resolution.
	/// Acts like indexed list of DeferRef
	struct DeferItem {
		DeclarationWithType* decl;
		AssertionSetValue info;
		CandidateList matches;

		DeferItem( DeclarationWithType* decl, const AssertionSetValue& info, 
			CandidateList&& matches )
		: decl(decl), info(info), matches(std::move(matches)) {}

		bool empty() const { return matches.empty(); }

		CandidateList::size_type size() const { return matches.size(); }

		DeferRef operator[] ( unsigned i ) const { return { decl, info, matches[i] }; }
	};

	/// List of deferred resolution items
	using DeferList = std::vector<DeferItem>;

	/// Combo iterator that combines candidates into an output list, merging their environments. 
	/// Rejects an appended candidate if the environments cannot be merged.
	class CandidateEnvMerger {
		/// Current list of merged candidates
		std::vector<DeferRef> crnt;
		/// Stack of environments to support backtracking
		std::vector<TypeEnvironment> envs;
		/// Stack of open variables to support backtracking
		std::vector<OpenVarSet> varSets;
		/// Indexer to use for merges
		const SymTab::Indexer& indexer;
	
	public:
		/// The merged environment/open variables and the list of candidates
		struct OutType {
			TypeEnvironment env;
			OpenVarSet openVars;
			std::vector<DeferRef> assns;

			OutType( const TypeEnvironment& env, const OpenVarSet& openVars, 
				const std::vector<DeferRef>& assns )
			: env(env), openVars(openVars), assns(assns) {}
		};

		CandidateEnvMerger( const TypeEnvironment& env, const OpenVarSet& openVars, 
			const SymTab::Indexer& indexer )
		: crnt(), envs{ env }, varSets{ openVars }, indexer(indexer) {}

		bool append( DeferRef i ) {
			TypeEnvironment env = envs.back();
			OpenVarSet openVars = varSets.back();
			mergeOpenVars( openVars, i.match.openVars );

			if ( ! env.combine( i.match.env, openVars, indexer ) ) return false;

			crnt.emplace_back( i );
			envs.emplace_back( env );
			varSets.emplace_back( openVars );
			return true;
		}

		void backtrack() {
			crnt.pop_back();
			envs.pop_back();
			varSets.pop_back();
		}

		OutType finalize() { return { envs.back(), varSets.back(), crnt }; }
	};

	/// Comparator for CandidateEnvMerger outputs that sums their costs and caches the stored 
	/// sums
	struct CandidateCost {
		using Element = CandidateEnvMerger::OutType;
	private:
		using Memo = std::unordered_map<const Element*, Cost>;
		mutable Memo cache;              ///< cache of element costs
		const SymTab::Indexer& indexer;  ///< Indexer for costing

	public:
		CandidateCost( const SymTab::Indexer& indexer ) : indexer(indexer) {}

		/// reports the cost of an element
		Cost get( const Element& x ) const {
			Memo::const_iterator it = cache.find( &x );
			if ( it == cache.end() ) {
				Cost k = Cost::zero;
				for ( const auto& assn : x.assns ) {
					k += computeConversionCost( 
						assn.match.adjType, assn.decl->get_type(), indexer, x.env );
				}
				it = cache.emplace_hint( it, &x, k );
			}
			return it->second;
		}
		
		/// compares elements by cost
		bool operator() ( const Element& a, const Element& b ) const {
			return get( a ) < get( b );
		}
	};

	/// Set of assertion resolutions, grouped by resolution ID
	using InferCache = std::unordered_map< UniqueId, InferredParams >;

	/// Flag for state iteration
	enum IterateFlag { IterateState };

	/// State needed to resolve a set of assertions
	struct ResnState {
		Alternative alt;           ///< Alternative assertion is rooted on
		AssertionList need;        ///< Assertions to find
		AssertionSet newNeed;      ///< New assertions for current resolutions
		DeferList deferred;        ///< Deferred matches
		InferCache inferred;       ///< Cache of already-inferred parameters
		SymTab::Indexer& indexer;  ///< Name lookup (depends on previous assertions)

		/// Initial resolution state for an alternative
		ResnState( Alternative& a, SymTab::Indexer& indexer )
		: alt(a), need(), newNeed(), deferred(), inferred(), indexer(indexer) {
			need.swap( a.need );
		}

		/// Updated resolution state with new need-list
		ResnState( ResnState&& o, IterateFlag )
		: alt(std::move(o.alt)), need(o.newNeed.begin(), o.newNeed.end()), newNeed(), deferred(), 
		  inferred(std::move(o.inferred)), indexer(o.indexer) {}
	};

	/// Binds a single assertion, updating resolution state
	void bindAssertion( const DeclarationWithType* decl, AssertionSetValue info, Alternative& alt, 
			AssnCandidate& match, InferCache& inferred ) {
		
		DeclarationWithType* candidate = match.cdata.id;
		assertf( candidate->get_uniqueId(), "Assertion candidate does not have a unique ID: %s", toString( candidate ).c_str() );

		Expression* varExpr = match.cdata.combine( alt.cvtCost );
		delete varExpr->result;
		varExpr->result = match.adjType->clone();
		if ( match.resnSlot ) { varExpr->resnSlots.push_back( match.resnSlot ); }

		// place newly-inferred assertion in proper place in cache
		inferred[ info.resnSlot ][ decl->get_uniqueId() ] = ParamEntry{
				candidate->get_uniqueId(), match.adjType->clone(), decl->get_type()->clone(), 
				varExpr };

		// // follow the current assertion's ID chain to find the correct set of inferred parameters 
		// // to add the candidate o (i.e. the set of inferred parameters belonging to the entity 
		// // which requested the assertion parameter)
		// InferredParams* inferParams = &alt.expr->inferParams;
		// for ( UniqueId id : info.idChain ) {
		// 	inferParams = (*inferParams)[ id ].inferParams.get();
		// }

		// (*inferParams)[ decl->get_uniqueId() ] = ParamEntry{
		// 		candidate->get_uniqueId(), match.adjType, decl->get_type()->clone(), varExpr };
	}

	/// Adds a captured assertion to the symbol table
	void addToIndexer( AssertionSet &assertSet, SymTab::Indexer &indexer ) {
		for ( AssertionSet::iterator i = assertSet.begin(); i != assertSet.end(); ++i ) {
			if ( i->second.isUsed ) {
				indexer.addId( i->first );
			}
		}
	}

	// in AlternativeFinder.cc; unique ID for assertion resolutions
	extern UniqueId globalResnSlot;

	/// Resolve a single assertion, in context
	bool resolveAssertion( AssertionItem& assn, ResnState& resn ) {
		// skip unused assertions
		if ( ! assn.info.isUsed ) return true;

		// lookup candidates for this assertion
		std::list< SymTab::Indexer::IdData > candidates;
		resn.indexer.lookupId( assn.decl->name, candidates );

		// find the candidates that unify with the desired type
		CandidateList matches;
		for ( const auto& cdata : candidates ) {
			DeclarationWithType* candidate = cdata.id;

			// build independent unification context for candidate
			AssertionSet have, newNeed;
			TypeEnvironment newEnv{ resn.alt.env };
			OpenVarSet newOpenVars{ resn.alt.openVars };
			Type* adjType = candidate->get_type()->clone();
			adjustExprType( adjType, newEnv, resn.indexer );
			renameTyVars( adjType );

			// keep unifying candidates
			if ( unify( assn.decl->get_type(), adjType, newEnv, newNeed, have, newOpenVars, 
					resn.indexer ) ) {
				// set up binding slot for recursive assertions
				UniqueId crntResnSlot = 0;
				if ( ! newNeed.empty() ) {
					crntResnSlot = ++globalResnSlot;
					for ( auto& a : newNeed ) {
						a.second.resnSlot = crntResnSlot;
					}
				}
				// // set up idChain on new assertions
				// for ( auto& a : newNeed ) {
				// 	a.second.idChain = assn.info.idChain;
				// 	a.second.idChain.push_back( assn.decl->get_uniqueId() );
				// }

				matches.emplace_back( cdata, adjType, std::move(newEnv), std::move(have), 
					std::move(newNeed), std::move(newOpenVars), crntResnSlot );
			} else {
				delete adjType;
			}
		}

		// break if no suitable assertion
		if ( matches.empty() ) return false;

		// defer if too many suitable assertions
		if ( matches.size() > 1 ) {
			resn.deferred.emplace_back( assn.decl, assn.info, std::move(matches) );
			return true;
		}

		// otherwise bind current match in ongoing scope
		AssnCandidate& match = matches.front();
		addToIndexer( match.have, resn.indexer );
		resn.newNeed.insert( match.need.begin(), match.need.end() );
		resn.alt.env = std::move(match.env);
		resn.alt.openVars = std::move(match.openVars);

		bindAssertion( assn.decl, assn.info, resn.alt, match, resn.inferred );
		return true;
	}

	/// Associates inferred parameters with an expression
	struct InferMatcher {
		InferCache& inferred;

		InferMatcher( InferCache& inferred ) : inferred( inferred ) {}

		Expression* postmutate( Expression* expr ) {
			// defer missing inferred parameters until they are hopefully found later
			std::vector<UniqueId> missingSlots;
			// place inferred parameters into resolution slots
			for ( UniqueId slot : expr->resnSlots ) {
				// fail if no matching parameters found
				auto it = inferred.find( slot );
				if ( it == inferred.end() ) {
					missingSlots.push_back( slot );
					continue;
				}
				InferredParams& inferParams = it->second;
				
				// place inferred parameters into proper place in expression
				for ( auto& entry : inferParams ) {
					// recurse on inferParams of resolved expressions
					entry.second.expr = postmutate( entry.second.expr );
					// xxx - look at entry.second.inferParams?
					expr->inferParams[ entry.first ] = entry.second;
				}
			}

			// clear resolution slots and return
			expr->resnSlots.swap( missingSlots );
			return expr;
		}
	};

	void finalizeAssertions( Alternative& alt, InferCache& inferred, AltList& out ) {
		PassVisitor<InferMatcher> matcher{ inferred };
		alt.expr = alt.expr->acceptMutator( matcher );
		out.emplace_back( alt );
	}

	/// Limit to depth of recursion of assertion satisfaction
	static const int recursionLimit = /* 10 */ 4;

	void resolveAssertions( Alternative& alt, const SymTab::Indexer& indexer, AltList& out ) {
		// finish early if no assertions to resolve
		if ( alt.need.empty() ) {
			out.emplace_back( alt );
			return;
		}

		// build list of possible resolutions
		using ResnList = std::vector<ResnState>;
		SymTab::Indexer root_indexer{ indexer };
		ResnList resns{ ResnState{ alt, root_indexer } };
		ResnList new_resns{};

		// resolve assertions in breadth-first-order up to a limited number of levels deep
		for ( unsigned level = 0; level < recursionLimit; ++level ) {
			// scan over all mutually-compatible resolutions
			for ( auto& resn : resns ) {
				// make initial pass at matching assertions
				for ( auto& assn : resn.need ) {
					// fail early if any assertion is not resolvable
					if ( ! resolveAssertion( assn, resn ) ) goto nextResn;
				}

				if ( resn.deferred.empty() ) {
					// either add successful match or push back next state
					if ( resn.newNeed.empty() ) {
						finalizeAssertions( resn.alt, resn.inferred, out );
					} else {
						new_resns.emplace_back( std::move(resn), IterateState );
					}
				} else {
					// resolve deferred assertions by mutual compatibility
					std::vector<CandidateEnvMerger::OutType> compatible = filterCombos(
						resn.deferred, 
						CandidateEnvMerger{ resn.alt.env, resn.alt.openVars, resn.indexer } );
					// sort by cost
					CandidateCost coster{ resn.indexer };
					std::sort( compatible.begin(), compatible.end(), coster );
					// // sort by cost if pruning
					// if ( pruneAssertions ) {
					// 	auto lmin = sort_mins( compatible.begin(), compatible.end(), 
					// 		CandidateCost{resn.indexer} );
					// 	compatible.erase( lmin, compatible.end() );
					// }

					// keep map of detected options
					std::unordered_map<std::string, Cost> found;
					for ( auto& compat : compatible ) {
						// filter by environment-adjusted result type, keep only cheapest option
						Type* resType = alt.expr->result->clone();
						compat.env.apply( resType );
						// skip if cheaper alternative already processed with same result type
						Cost resCost = coster.get( compat );
						auto it = found.emplace( SymTab::Mangler::mangleType( resType ), resCost );
						delete resType;
						if ( it.second == false && it.first->second < resCost ) continue;

						// proceed with resolution step
						ResnState new_resn = resn;

						// add compatible assertions to new resolution state
						for ( DeferRef r : compat.assns ) {
							AssnCandidate match = r.match;
							addToIndexer( match.have, new_resn.indexer );
							new_resn.newNeed.insert( match.need.begin(), match.need.end() );

							bindAssertion( r.decl, r.info, new_resn.alt, match, new_resn.inferred );
						}

						// set mutual environment into resolution state
						new_resn.alt.env = std::move(compat.env);
						new_resn.alt.openVars = std::move(compat.openVars);

						// either add sucessful match or push back next state
						if ( new_resn.newNeed.empty() ) {
							finalizeAssertions( new_resn.alt, new_resn.inferred, out );
						} else {
							new_resns.emplace_back( std::move(new_resn), IterateState );
						}
					}
				}
			nextResn:; }

			// finish or reset for next round
			if ( new_resns.empty() ) return;
			resns.swap( new_resns );
			new_resns.clear();
		}
		
		// exceeded recursion limit if reaches here
		if ( out.empty() ) {
			SemanticError( alt.expr->location, "Too many recursive assertions" );
		}
	}
} // namespace ResolvExpr

// Local Variables: //
// tab-width: 4 //
// mode: c++ //
// compile-command: "make install" //
// End: //