source: src/Tuples/TupleExpansionNew.cpp @ e116db3

//
// 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.
//
// TupleExpansionNew.cpp --
//
// Author           : Henry Xue
// Created On       : Mon Aug 23 15:36:09 2021
// Last Modified By : Andrew Beach
// Last Modified On : Mon Aug 15 17:00:00 2022
// Update Count     : 3
//

#include "Tuples.h"

#include "AST/Pass.hpp"
#include "Common/ScopedMap.h"

namespace Tuples {

namespace {

struct MemberTupleExpander final : public ast::WithShortCircuiting, public ast::WithVisitorRef< MemberTupleExpander > {
        void previsit( const ast::UntypedMemberExpr * ) { visit_children = false; }
        const ast::Expr * postvisit( const ast::UntypedMemberExpr * memberExpr );
};

struct UniqueExprExpander final : public ast::WithDeclsToAdd<> {
        const ast::Expr * postvisit( const ast::UniqueExpr * unqExpr );
        std::map< int, ast::ptr<ast::Expr> > decls; // not vector, because order added may not be increasing order
};

/// given a expression representing the member and an expression representing the aggregate,
/// reconstructs a flattened UntypedMemberExpr with the right precedence
const ast::Expr * reconstructMemberExpr( const ast::Expr * member, const ast::Expr * aggr, const CodeLocation & loc ) {
        if ( auto memberExpr = dynamic_cast< const ast::UntypedMemberExpr * >( member ) ) {
                // construct a new UntypedMemberExpr with the correct structure , and recursively
                // expand that member expression.
                ast::Pass< MemberTupleExpander > expander;
                auto inner = new ast::UntypedMemberExpr( loc, memberExpr->aggregate, aggr );
                auto newMemberExpr = new ast::UntypedMemberExpr( loc, memberExpr->member, inner );
                return newMemberExpr->accept( expander );
        } else {
                // not a member expression, so there is nothing to do but attach and return
                return new ast::UntypedMemberExpr( loc, member, aggr );
        }
}

const ast::Expr * MemberTupleExpander::postvisit( const ast::UntypedMemberExpr * memberExpr ) {
        const CodeLocation loc = memberExpr->location;
        if ( auto tupleExpr = memberExpr->member.as< ast::UntypedTupleExpr >() ) {
                auto mutExpr = mutate( tupleExpr );
                ast::ptr< ast::Expr > aggr = memberExpr->aggregate->accept( *visitor );
                // aggregate expressions which might be impure must be wrapped in unique expressions
                if ( Tuples::maybeImpureIgnoreUnique( memberExpr->aggregate ) ) aggr = new ast::UniqueExpr( loc, aggr );
                for ( auto & expr : mutExpr->exprs ) {
                        expr = reconstructMemberExpr( expr, aggr, loc );
                }
                return mutExpr;
        } else {
                // there may be a tuple expr buried in the aggregate
                return new ast::UntypedMemberExpr( loc, memberExpr->member, memberExpr->aggregate->accept( *visitor ) );
        }
}

const ast::Expr * UniqueExprExpander::postvisit( const ast::UniqueExpr * unqExpr ) {
        const CodeLocation loc = unqExpr->location;
        const int id = unqExpr->id;

        // on first time visiting a unique expr with a particular ID, generate the expression that replaces all UniqueExprs with that ID,
        // and lookup on subsequent hits. This ensures that all unique exprs with the same ID reference the same variable.
        if ( ! decls.count( id ) ) {
                ast::ptr< ast::Expr > assignUnq;
                const ast::VariableExpr * var = unqExpr->var;
                if ( unqExpr->object ) {
                        // an object was generated to represent this unique expression -- it should be added to the list of declarations now
                        declsToAddBefore.push_back( unqExpr->object.as< ast::Decl >() );
                        // deep copy required due to unresolved issues with UniqueExpr
                        assignUnq = ast::UntypedExpr::createAssign( loc, var, unqExpr->expr );
                } else {
                        const auto commaExpr = unqExpr->expr.strict_as< ast::CommaExpr >();
                        assignUnq = commaExpr->arg1;
                }
                auto finished = new ast::ObjectDecl( loc, toString( "_unq", id, "_finished_" ), new ast::BasicType( ast::BasicType::Kind::Bool ),
                        new ast::SingleInit( loc, ast::ConstantExpr::from_int( loc, 0 ) ), {}, ast::Linkage::Cforall );
                declsToAddBefore.push_back( finished );
                // (finished ? _unq_expr_N : (_unq_expr_N = <unqExpr->get_expr()>, finished = 1, _unq_expr_N))
                // This pattern ensures that each unique expression is evaluated once, regardless of evaluation order of the generated C code.
                auto assignFinished = ast::UntypedExpr::createAssign( loc, new ast::VariableExpr( loc, finished ),
                        ast::ConstantExpr::from_int( loc, 1 ) );
                auto condExpr = new ast::ConditionalExpr( loc, new ast::VariableExpr( loc, finished ), var,
                        new ast::CommaExpr( loc, new ast::CommaExpr( loc, assignUnq, assignFinished ), var ) );
                condExpr->result = var->result;
                condExpr->env = unqExpr->env;
                decls[id] = condExpr;
        }
        return ast::deepCopy(decls[id].get());
}

/// Replaces Tuple Assign & Index Expressions, and Tuple Types.
struct TupleMainExpander :
                public ast::WithGuards,
                public ast::WithVisitorRef<TupleMainExpander>,
                public ast::WithDeclsToAdd<> {
        ast::Expr const * postvisit( ast::TupleAssignExpr const * expr ) {
                // Just move the env on the new top level expression.
                return ast::mutate_field( expr->stmtExpr.get(),
                        &ast::TupleAssignExpr::env, expr->env.get() );
        }

        void previsit( ast::ParseNode const * node ) {
                GuardValue( location ) = &node->location;
        }

        void previsit( ast::CompoundStmt const * stmt ) {
                previsit( (ast::ParseNode const *)stmt );
                GuardScope( typeMap );
        }

        ast::Expr const * postvisit( ast::Expr const * expr ) {
                if ( nullptr == expr->env ) {
                        return expr;
                }
                // TypeSubstitutions are never visited in the new Pass template,
                // so it is done manually here, where all types have to be replaced.
                return ast::mutate_field( expr, &ast::Expr::env,
                        expr->env->accept( *visitor ) );
        }

        ast::Type const * postvisit( ast::TupleType const * type ) {
                assert( location );
                unsigned tupleSize = type->size();
                if ( !typeMap.count( tupleSize ) ) {
                        ast::StructDecl * decl = new ast::StructDecl( *location,
                                toString( "_tuple", tupleSize, "_" )
                        );
                        decl->body = true;

                        for ( size_t i = 0 ; i < tupleSize ; ++i ) {
                                ast::TypeDecl * typeParam = new ast::TypeDecl( *location,
                                        toString( "tuple_param_", tupleSize, "_", i ),
                                        ast::Storage::Classes(),
                                        nullptr,
                                        ast::TypeDecl::Dtype,
                                        true
                                        );
                                ast::ObjectDecl * member = new ast::ObjectDecl( *location,
                                        toString( "field_", i ),
                                        new ast::TypeInstType( typeParam )
                                        );
                                decl->params.push_back( typeParam );
                                decl->members.push_back( member );
                        }

                        // Empty structures are not standard C. Add a dummy field to
                        // empty tuples to silence warnings when a compound literal
                        // `_tuple0_` is created.
                        if ( tupleSize == 0 ) {
                                decl->members.push_back(
                                        new ast::ObjectDecl( *location,
                                                "dummy",
                                                new ast::BasicType( ast::BasicType::SignedInt ),
                                                nullptr,
                                                ast::Storage::Classes(),
                                                // Does this have to be a C linkage?
                                                ast::Linkage::C
                                        )
                                );
                        }
                        typeMap[tupleSize] = decl;
                        declsToAddBefore.push_back( decl );
                }

                ast::StructDecl const * decl = typeMap[ tupleSize ];
                ast::StructInstType * newType =
                        new ast::StructInstType( decl, type->qualifiers );
                for ( auto pair : group_iterate( type->types, decl->params ) ) {
                        ast::Type const * t = std::get<0>( pair );
                        newType->params.push_back(
                                new ast::TypeExpr( *location, ast::deepCopy( t ) ) );
                }
                return newType;
        }

        ast::Expr const * postvisit( ast::TupleIndexExpr const * expr ) {
                CodeLocation const & location = expr->location;
                ast::Expr const * tuple = expr->tuple.get();
                assert( tuple );
                unsigned int index = expr->index;
                ast::TypeSubstitution const * env = expr->env.get();

                if ( auto tupleExpr = dynamic_cast<ast::TupleExpr const *>( tuple ) ) {
                        // Optimization: If it is a definitely pure tuple expr,
                        // then it can reduce to the only relevant component.
                        if ( !maybeImpureIgnoreUnique( tupleExpr ) ) {
                                assert( index < tupleExpr->exprs.size() );
                                ast::ptr<ast::Expr> const & expr =
                                        *std::next( tupleExpr->exprs.begin(), index );
                                ast::Expr * ret = ast::mutate( expr.get() );
                                ret->env = env;
                                return ret;
                        }
                }

                auto type = tuple->result.strict_as<ast::StructInstType>();
                ast::StructDecl const * structDecl = type->base;
                assert( index < structDecl->members.size() );
                ast::ptr<ast::Decl> const & member =
                        *std::next( structDecl->members.begin(), index );
                ast::MemberExpr * memberExpr = new ast::MemberExpr( location,
                        member.strict_as<ast::DeclWithType>(), tuple );
                memberExpr->env = env;
                return memberExpr;
        }
private:
        ScopedMap< int, ast::StructDecl const * > typeMap;
        CodeLocation const * location = nullptr;
};

ast::Expr const * replaceTupleExpr(
                CodeLocation const & location,
                ast::Type const * result,
                std::vector<ast::ptr<ast::Expr>> const & exprs,
                ast::TypeSubstitution const * env ) {
        assert( result );
        // A void result: It doesn't need to produce a value for cascading,
        // just output a chain of comma exprs.
        if ( result->isVoid() ) {
                assert( !exprs.empty() );
                std::vector<ast::ptr<ast::Expr>>::const_iterator iter = exprs.begin();
                ast::Expr * expr = new ast::CastExpr( *iter++ );
                for ( ; iter != exprs.end() ; ++iter ) {
                        expr = new ast::CommaExpr( location,
                                expr, new ast::CastExpr( *iter ) );
                }
                expr->env = env;
                return expr;
        // Typed tuple expression: Produce a compound literal which performs
        // each of the expressions as a distinct part of its initializer. The
        // produced compound literal may be used as part of another expression.
        } else {
                auto inits = map_range<std::vector<ast::ptr<ast::Init>>>( exprs,
                        []( ast::Expr const * expr ) {
                                return new ast::SingleInit( expr->location, expr );
                        }
                );
                ast::Expr * expr = new ast::CompoundLiteralExpr( location,
                        result, new ast::ListInit( location, std::move( inits ) ) );
                expr->env = env;
                return expr;
        }
}

struct TupleExprExpander {
        ast::Expr const * postvisit( ast::TupleExpr const * expr ) {
                return replaceTupleExpr( expr->location,
                        expr->result, expr->exprs, expr->env );
        }
};

} // namespace

void expandMemberTuples( ast::TranslationUnit & translationUnit ) {
        ast::Pass< MemberTupleExpander >::run( translationUnit );
}

void expandUniqueExpr( ast::TranslationUnit & translationUnit ) {
        ast::Pass< UniqueExprExpander >::run( translationUnit );
}

void expandTuples( ast::TranslationUnit & translationUnit ) {
        // These can't just be combined simply (there might be a way with work).
        ast::Pass<TupleMainExpander>::run( translationUnit );
        ast::Pass<TupleExprExpander>::run( translationUnit );
}

} // namespace Tuples