source: src/GenPoly/SpecializeNew.cpp @ 0e8f4c6

//
// 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.
//
// SpecializeNew.cpp -- Generate thunks to specialize polymorphic functions.
//
// Author           : Andrew Beach
// Created On       : Tue Jun  7 13:37:00 2022
// Last Modified By : Andrew Beach
// Last Modified On : Tue Jun  7 13:37:00 2022
// Update Count     : 0
//

#include "Specialize.h"

#include "AST/Copy.hpp"                  // for deepCopy
#include "AST/Inspect.hpp"               // for isIntrinsicCallExpr
#include "AST/Pass.hpp"                  // for Pass
#include "AST/TypeEnvironment.hpp"       // for OpenVarSet, AssertionSet
#include "Common/UniqueName.h"           // for UniqueName
#include "GenPoly/GenPoly.h"             // for getFunctionType
#include "ResolvExpr/FindOpenVars.h"     // for findOpenVars
#include "ResolvExpr/TypeEnvironment.h"  // for FirstOpen, FirstClosed

namespace GenPoly {

namespace {

struct SpecializeCore final :
                public ast::WithConstTypeSubstitution,
                public ast::WithDeclsToAdd<>,
                public ast::WithVisitorRef<SpecializeCore> {
        std::string paramPrefix = "_p";

        ast::ApplicationExpr * handleExplicitParams(
                const ast::ApplicationExpr * expr );
        const ast::Expr * createThunkFunction(
                const CodeLocation & location,
                const ast::FunctionType * funType,
                const ast::Expr * actual,
                const ast::InferredParams * inferParams );
        const ast::Expr * doSpecialization(
                const CodeLocation & location,
                const ast::Type * formalType,
                const ast::Expr * actual,
                const ast::InferredParams * inferParams );

        const ast::Expr * postvisit( const ast::ApplicationExpr * expr );
        const ast::Expr * postvisit( const ast::CastExpr * expr );
};

const ast::InferredParams * getInferredParams( const ast::Expr * expr ) {
        const ast::Expr::InferUnion & inferred = expr->inferred;
        if ( inferred.hasParams() ) {
                return &inferred.inferParams();
        } else {
                return nullptr;
        }
}

// Check if both types have the same structure. The leaf (non-tuple) types
// don't have to match but the tuples must match.
bool isTupleStructureMatching( const ast::Type * t0, const ast::Type * t1 ) {
        const ast::TupleType * tt0 = dynamic_cast<const ast::TupleType *>( t0 );
        const ast::TupleType * tt1 = dynamic_cast<const ast::TupleType *>( t1 );
        if ( tt0 && tt1 ) {
                if ( tt0->size() != tt1->size() ) {
                        return false;
                }
                for ( auto types : group_iterate( tt0->types, tt1->types ) ) {
                        if ( !isTupleStructureMatching(
                                        std::get<0>( types ), std::get<1>( types ) ) ) {
                                return false;
                        }
                }
                return true;
        }
        return (!tt0 && !tt1);
}

// The number of elements in a type if it is a flattened tuple.
size_t flatTupleSize( const ast::Type * type ) {
        if ( auto tuple = dynamic_cast<const ast::TupleType *>( type ) ) {
                size_t sum = 0;
                for ( auto t : *tuple ) {
                        sum += flatTupleSize( t );
                }
                return sum;
        } else {
                return 1;
        }
}

// Find the total number of components in a parameter list.
size_t functionParameterSize( const ast::FunctionType * type ) {
        size_t sum = 0;
        for ( auto param : type->params ) {
                sum += flatTupleSize( param );
        }
        return sum;
}

bool needsPolySpecialization(
                const ast::Type * /*formalType*/,
                const ast::Type * actualType,
                const ast::TypeSubstitution * subs ) {
        if ( !subs ) {
                return false;
        }

        using namespace ResolvExpr;
        ast::OpenVarSet openVars, closedVars;
        ast::AssertionSet need, have; // unused
        ast::TypeEnvironment env; // unused
        // findOpenVars( formalType, openVars, closedVars, need, have, FirstClosed );
        findOpenVars( actualType, openVars, closedVars, need, have, env, FirstOpen );
        for ( const ast::OpenVarSet::value_type & openVar : openVars ) {
                const ast::Type * boundType = subs->lookup( openVar.first );
                // If the variable is not bound, move onto the next variable.
                if ( !boundType ) continue;

                // Is the variable cound to another type variable?
                if ( auto inst = dynamic_cast<const ast::TypeInstType *>( boundType ) ) {
                        if ( closedVars.find( *inst ) == closedVars.end() ) {
                                return true;
                        } else {
                                assertf(false, "closed: %s", inst->name.c_str());
                        }
                // Otherwise, the variable is bound to a concrete type.
                } else {
                        return true;
                }
        }
        // None of the type variables are bound.
        return false;
}

bool needsTupleSpecialization(
                const ast::Type * formalType, const ast::Type * actualType ) {
        // Needs tuple specialization if the structure of the formal type and
        // actual type do not match.

        // This is the case if the formal type has ttype polymorphism, or if the structure  of tuple types
        // between the function do not match exactly.
        if ( const ast::FunctionType * ftype = getFunctionType( formalType ) ) {
                // A pack in the parameter or return type requires specialization.
                if ( ftype->isTtype() ) {
                        return true;
                }
                // Conversion of 0 to a function type does not require specialization.
                if ( dynamic_cast<const ast::ZeroType *>( actualType ) ) {
                        return false;
                }
                const ast::FunctionType * atype =
                        getFunctionType( actualType->stripReferences() );
                assertf( atype,
                        "formal type is a function type, but actual type is not: %s",
                        toString( actualType ).c_str() );
                // Can't tuple specialize if parameter sizes deeply-differ.
                if ( functionParameterSize( ftype ) != functionParameterSize( atype ) ) {
                        return false;
                }
                // If tuple parameter size matches but actual parameter sizes differ
                // then there needs to be specialization.
                if ( ftype->params.size() != atype->params.size() ) {
                        return true;
                }
                // Total parameter size can be the same, while individual parameters
                // can have different structure.
                for ( auto pairs : group_iterate( ftype->params, atype->params ) ) {
                        if ( !isTupleStructureMatching(
                                        std::get<0>( pairs ), std::get<1>( pairs ) ) ) {
                                return true;
                        }
                }
        }
        return false;
}

bool needsSpecialization(
                const ast::Type * formalType, const ast::Type * actualType,
                const ast::TypeSubstitution * subs ) {
        return needsPolySpecialization( formalType, actualType, subs )
                || needsTupleSpecialization( formalType, actualType );
}

ast::ApplicationExpr * SpecializeCore::handleExplicitParams(
                const ast::ApplicationExpr * expr ) {
        assert( expr->func->result );
        const ast::FunctionType * func = getFunctionType( expr->func->result );
        assert( func );

        ast::ApplicationExpr * mut = ast::mutate( expr );

        std::vector<ast::ptr<ast::Type>>::const_iterator formal;
        std::vector<ast::ptr<ast::Expr>>::iterator actual;
        for ( formal = func->params.begin(), actual = mut->args.begin() ;
                        formal != func->params.end() && actual != mut->args.end() ;
                        ++formal, ++actual ) {
                *actual = doSpecialization( (*actual)->location,
                        *formal, *actual, getInferredParams( expr ) );
        }
        return mut;
}

// Explode assuming simple cases: either type is pure tuple (but not tuple
// expr) or type is non-tuple.
template<typename OutputIterator>
void explodeSimple( const CodeLocation & location,
                const ast::Expr * expr, OutputIterator out ) {
        // Recurse on tuple types using index expressions on each component.
        if ( auto tuple = expr->result.as<ast::TupleType>() ) {
                ast::ptr<ast::Expr> cleanup = expr;
                for ( unsigned int i = 0 ; i < tuple->size() ; ++i ) {
                        explodeSimple( location,
                                new ast::TupleIndexExpr( location, expr, i ), out );
                }
        // For a non-tuple type, output a clone of the expression.
        } else {
                *out++ = expr;
        }
}

// Restructures arguments to match the structure of the formal parameters
// of the actual function. Returns the next structured argument.
template<typename Iterator>
const ast::Expr * structureArg(
                const CodeLocation& location, const ast::ptr<ast::Type> & type,
                Iterator & begin, const Iterator & end ) {
        if ( auto tuple = type.as<ast::TupleType>() ) {
                std::vector<ast::ptr<ast::Expr>> exprs;
                for ( const ast::ptr<ast::Type> & t : *tuple ) {
                        exprs.push_back( structureArg( location, t, begin, end ) );
                }
                return new ast::TupleExpr( location, std::move( exprs ) );
        } else {
                assertf( begin != end, "reached the end of the arguments while structuring" );
                return *begin++;
        }
}

struct TypeInstFixer final : public ast::WithShortCircuiting {
        std::map<const ast::TypeDecl *, std::pair<int, int>> typeMap;

        void previsit(const ast::TypeDecl *) { visit_children = false; }
        const ast::TypeInstType * postvisit(const ast::TypeInstType * typeInst) {
                if (typeMap.count(typeInst->base)) {
                        ast::TypeInstType * newInst = mutate(typeInst);
                        auto const & pair = typeMap[typeInst->base];
                        newInst->expr_id = pair.first;
                        newInst->formal_usage = pair.second;
                        return newInst;
                }
                return typeInst;
        }
};

const ast::Expr * SpecializeCore::createThunkFunction(
                const CodeLocation & location,
                const ast::FunctionType * funType,
                const ast::Expr * actual,
                const ast::InferredParams * inferParams ) {
        // One set of unique names per program.
        static UniqueName thunkNamer("_thunk");

        const ast::FunctionType * newType = ast::deepCopy( funType );
        if ( typeSubs ) {
                // Must replace only occurrences of type variables
                // that occure free in the thunk's type.
                auto result = typeSubs->applyFree( newType );
                newType = result.node.release();
        }

        using DWTVector = std::vector<ast::ptr<ast::DeclWithType>>;
        using DeclVector = std::vector<ast::ptr<ast::TypeDecl>>;

        UniqueName paramNamer( paramPrefix );

        // Create new thunk with same signature as formal type.
        ast::Pass<TypeInstFixer> fixer;
        for (const auto & kv : newType->forall) {
                if (fixer.core.typeMap.count(kv->base)) {
                        std::cerr << location << ' ' << kv->base->name
                                << ' ' << kv->expr_id << '_' << kv->formal_usage
                                << ',' << fixer.core.typeMap[kv->base].first
                                << '_' << fixer.core.typeMap[kv->base].second << std::endl;
                        assertf(false, "multiple formals in specialize");
                }
                else {
                        fixer.core.typeMap[kv->base] = std::make_pair(kv->expr_id, kv->formal_usage);
                }
        }

        ast::CompoundStmt * thunkBody = new ast::CompoundStmt( location );
        ast::FunctionDecl * thunkFunc = new ast::FunctionDecl(
                location,
                thunkNamer.newName(),
                map_range<DeclVector>( newType->forall, []( const ast::TypeInstType * inst ) {
                        return ast::deepCopy( inst->base );
                } ),
                map_range<DWTVector>( newType->assertions, []( const ast::VariableExpr * expr ) {
                        return ast::deepCopy( expr->var );
                } ),
                map_range<DWTVector>( newType->params, [&location, &paramNamer]( const ast::Type * type ) {
                        return new ast::ObjectDecl( location, paramNamer.newName(), ast::deepCopy( type ) );
                } ),
                map_range<DWTVector>( newType->returns, [&location, &paramNamer]( const ast::Type * type ) {
                        return new ast::ObjectDecl( location, paramNamer.newName(), ast::deepCopy( type ) );
                } ),
                thunkBody,
                ast::Storage::Classes(),
                ast::Linkage::C
                );

        thunkFunc->fixUniqueId();

        // Thunks may be generated and not used, avoid them.
        thunkFunc->attributes.push_back( new ast::Attribute( "unused" ) );

        // Global thunks must be static to avoid collitions.
        // Nested thunks must not be unique and hence, not static.
        thunkFunc->storage.is_static = !isInFunction();

        // Weave thunk parameters into call to actual function,
        // naming thunk parameters as we go.
        ast::ApplicationExpr * app = new ast::ApplicationExpr( location, actual );

        const ast::FunctionType * actualType = ast::deepCopy( getFunctionType( actual->result ) );
        if ( typeSubs ) {
                // Need to apply the environment to the actual function's type,
                // since it may itself be polymorphic.
                auto result = typeSubs->apply( actualType );
                actualType = result.node.release();
        }

        ast::ptr<ast::FunctionType> actualTypeManager = actualType;

        std::vector<ast::ptr<ast::Expr>> args;
        for ( ast::ptr<ast::DeclWithType> & param : thunkFunc->params ) {
                // Name each thunk parameter and explode it.
                // These are then threaded back into the actual function call.
                ast::DeclWithType * mutParam = ast::mutate( param.get() );
                explodeSimple( location, new ast::VariableExpr( location, mutParam ),
                        std::back_inserter( args ) );
        }

        // Walk parameters to the actual function alongside the exploded thunk
        // parameters and restructure the arguments to match the actual parameters.
        std::vector<ast::ptr<ast::Expr>>::iterator
                argBegin = args.begin(), argEnd = args.end();
        for ( const auto & actualArg : actualType->params ) {
                app->args.push_back(
                        structureArg( location, actualArg.get(), argBegin, argEnd ) );
        }
        assertf( argBegin == argEnd, "Did not structure all arguments." );

        app->accept(fixer); // this should modify in place

        app->env = ast::TypeSubstitution::newFromExpr( app, typeSubs );
        if ( inferParams ) {
                app->inferred.inferParams() = *inferParams;
        }

        // Handle any specializations that may still be present.
        {
                std::string oldParamPrefix = paramPrefix;
                paramPrefix += "p";
                std::list<ast::ptr<ast::Decl>> oldDecls;
                oldDecls.splice( oldDecls.end(), declsToAddBefore );

                app->accept( *visitor );
                // Write recursive specializations into the thunk body.
                for ( const ast::ptr<ast::Decl> & decl : declsToAddBefore ) {
                        thunkBody->push_back( new ast::DeclStmt( decl->location, decl ) );
                }

                declsToAddBefore = std::move( oldDecls );
                paramPrefix = std::move( oldParamPrefix );
        }

        // Add return (or valueless expression) to the thunk.
        ast::Stmt * appStmt;
        if ( funType->returns.empty() ) {
                appStmt = new ast::ExprStmt( app->location, app );
        } else {
                appStmt = new ast::ReturnStmt( app->location, app );
        }
        thunkBody->push_back( appStmt );

        // Add the thunk definition:
        declsToAddBefore.push_back( thunkFunc );

        // Return address of thunk function as replacement expression.
        return new ast::AddressExpr( location,
                new ast::VariableExpr( location, thunkFunc ) );
}

const ast::Expr * SpecializeCore::doSpecialization(
                const CodeLocation & location,
                const ast::Type * formalType,
                const ast::Expr * actual,
                const ast::InferredParams * inferParams ) {
        assertf( actual->result, "attempting to specialize an untyped expression" );
        if ( needsSpecialization( formalType, actual->result, typeSubs ) ) {
                if ( const ast::FunctionType * type = getFunctionType( formalType ) ) {
                        if ( const ast::ApplicationExpr * expr =
                                        dynamic_cast<const ast::ApplicationExpr *>( actual ) ) {
                                return createThunkFunction( location, type, expr->func, inferParams );
                        } else if ( auto expr =
                                        dynamic_cast<const ast::VariableExpr *>( actual ) ) {
                                return createThunkFunction( location, type, expr, inferParams );
                        } else {
                                // (I don't even know what that comment means.)
                                // This likely won't work, as anything that could build an ApplicationExpr probably hit one of the previous two branches
                                return createThunkFunction( location, type, actual, inferParams );
                        }
                } else {
                        return actual;
                }
        } else {
                return actual;
        }
}

const ast::Expr * SpecializeCore::postvisit(
                const ast::ApplicationExpr * expr ) {
        if ( ast::isIntrinsicCallExpr( expr ) ) {
                return expr;
        }

        // Create thunks for the inferred parameters.
        // This is not needed for intrinsic calls, because they aren't
        // actually passed to the function. It needs to handle explicit params
        // before inferred params so that explicit params do not recieve a
        // changed set of inferParams (and change them again).
        // Alternatively, if order starts to matter then copy expr's inferParams
        // and pass them to handleExplicitParams.
        ast::ApplicationExpr * mut = handleExplicitParams( expr );
        if ( !mut->inferred.hasParams() ) {
                return mut;
        }
        ast::InferredParams & inferParams = mut->inferred.inferParams();
        for ( ast::InferredParams::value_type & inferParam : inferParams ) {
                inferParam.second.expr = doSpecialization(
                        inferParam.second.expr->location,
                        inferParam.second.formalType,
                        inferParam.second.expr,
                        getInferredParams( inferParam.second.expr )
                );
        }
        return mut;
}

const ast::Expr * SpecializeCore::postvisit( const ast::CastExpr * expr ) {
        if ( expr->result->isVoid() ) {
                // No specialization if there is no return value.
                return expr;
        }
        const ast::Expr * specialized = doSpecialization(
                expr->location, expr->result, expr->arg, getInferredParams( expr ) );
        if ( specialized != expr->arg ) {
                // Assume that the specialization incorporates the cast.
                return specialized;
        } else {
                return expr;
        }
}

} // namespace

void convertSpecializations( ast::TranslationUnit & translationUnit ) {
        ast::Pass<SpecializeCore>::run( translationUnit );
}

} // namespace GenPoly

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