source: src/ControlStruct/TranslateEnumRange.cpp@ 7552fde

#include "TranslateEnumRange.hpp"

#include "AST/Pass.hpp"
#include "AST/TranslationUnit.hpp"

namespace ControlStruct {

struct addInitType {
    const ast::Stmt * postvisit( const ast::ForStmt * stmt );
};

struct addInit {
    const ast::Stmt * postvisit( const ast::ForStmt * stmt );
};

struct translateEnumRangeCore {
    const ast::Stmt * postvisit( const ast::ForStmt * stmt );
};

const ast::Stmt* addInitType::postvisit( const ast::ForStmt * stmt ) {
    if ( stmt->range_over ) {
        auto typeExpr = stmt->range_over.strict_as<ast::TypeExpr>();
        auto type = typeExpr->type;
        auto inits = stmt->inits;
        auto enumInst = type.strict_as<ast::EnumInstType>();
        auto enumDecl = enumInst->base;

        auto objInit = stmt->inits.front();

        auto initLocation = objInit->location;
        auto rangeLocation = stmt->range_over->location;
        assert( stmt->inits.size() == 1 );

        if (auto declStmt = objInit.as<ast::DeclStmt>()) {
            auto decl = declStmt->decl;
            if ( auto objDecl = decl.as<ast::ObjectDecl>()) {
                if ( !objDecl->type && type ) {
                    auto objDeclWithType = ast::mutate_field( objDecl, &ast::ObjectDecl::type, type );
                    auto declWithType = ast::mutate_field( declStmt, &ast::DeclStmt::decl, objDeclWithType );
                    stmt = ast::mutate_field_index( stmt, &ast::ForStmt::inits, 0, declWithType );
                }
            }
        }
    }
    return stmt;
}

const ast::Stmt* addInit::postvisit( const ast::ForStmt * stmt ) {
    if ( stmt->range_over ) {
        auto typeExpr = stmt->range_over.strict_as<ast::TypeExpr>();
        auto type = typeExpr->type;
        auto inits = stmt->inits;
        auto enumInst = type.strict_as<ast::EnumInstType>();
        auto enumDecl = enumInst->base;

        auto init = stmt->inits.front();

        if (auto declStmt = init.as<ast::DeclStmt>()) {
            auto decl = declStmt->decl;
            if ( auto objDecl = decl.as<ast::ObjectDecl>()) {
                if ( !objDecl->init ) {
                    auto location = stmt->location;
                    ast::SingleInit * newInit;
                    newInit = new ast::SingleInit( location, 
                        new ast::NameExpr( location, 
                        stmt->is_inc? enumDecl->members.front()->name: enumDecl->members.back()->name ) );
                    auto objDeclWithInit = ast::mutate_field( objDecl, &ast::ObjectDecl::init, newInit );
                    auto declWithInit = ast::mutate_field( declStmt, &ast::DeclStmt::decl, objDeclWithInit );
                    stmt = ast::mutate_field_index( stmt, &ast::ForStmt::inits, 0, declWithInit );
                }
            }
        }
    }
    return stmt;
}

const ast::Stmt* translateEnumRangeCore::postvisit( const ast::ForStmt * stmt ) {
    if ( !stmt->range_over ) return stmt;
    auto location = stmt->location;
    auto declStmt = stmt->inits.front().strict_as<ast::DeclStmt>();
    auto initDecl = declStmt->decl.strict_as<ast::ObjectDecl>();
    auto indexName = initDecl->name;
    
    auto typeExpr = stmt->range_over.strict_as<ast::TypeExpr>();
    auto enumInst = typeExpr->type.strict_as<ast::EnumInstType>();
    auto enumDecl = enumInst->base;
    
    // Both inc and dec check if the current posn less than the number of enumerator
    // for dec, it keeps call pred until it passes 0 (the first enumerator) and underflow,
    // it wraps around and become unsigned max
    ast::UntypedExpr * condition = ast::UntypedExpr::createCall( location,
        "?<=?",
        { 
            ast::UntypedExpr::createCall(location, "posn", { new ast::NameExpr( location, indexName ) } ),
            ast::ConstantExpr::from_ulong( location, enumDecl->members.size()-1 )
        });
    auto increment = ast::UntypedExpr::createCall( location, 
        stmt->is_inc? "succ": "pred",
        { new ast::NameExpr( location, indexName ) });
    auto assig = ast::UntypedExpr::createAssign( location, new ast::NameExpr( location, indexName ), increment );
    auto mut = ast::mutate_field( stmt, &ast::ForStmt::cond, condition );
    mut = ast::mutate_field(stmt, &ast::ForStmt::inc, assig );
    return mut;
}

void translateEnumRange( ast::TranslationUnit & translationUnit ) {
    ast::Pass<addInitType>::run( translationUnit );
    ast::Pass<addInit>::run( translationUnit );
    ast::Pass<translateEnumRangeCore>::run( translationUnit );
}
}