module sema.ScopeBuilder;

import tango.io.Stdout,
       tango.core.Array : find;

public
import sema.Scope,
       sema.Symbol;

import sema.Visitor,
       basic.SmallArray;

class ForwardReference : Visitor!(void)
{
    override void visit(Module[] modules)
    {
        (new TypeBuilder).visit(modules);
        this.modules = modules;
        inFunctionBodyStack.push(false);
        foreach (mod; modules)
        {
            current = mod;
            foreach (decl; mod.decls)
                visitDecl(decl);
        }
    }

    override void visitFuncDecl(FuncDecl d)
    {

        visitExp(d.returnType);
        visitExp(d.identifier);

        d.sym = current.symbol.createMember(
                d.identifier.get, 
                d.type, 
                d);

        auto old = current.symbol;
        current.symbol = d.sym;

        foreach (arg; d.funcArgs)
            visitDecl(arg);

        inFunctionBodyStack.push(true);

        foreach (stmt; d.statements)
            visitStmt(stmt);

        inFunctionBodyStack.pop();

        current.symbol = old;

    }

    override void visitVarDecl(VarDecl d)
    {
        visitExp(d.varType);
        visitExp(d.identifier);
        
        if (d.init)
            visitExp(d.init);

        DType t = typeOf(d.varType, d.env);
        d.sym = current.symbol.createAlias(
                d.identifier.get,
                d.env.find(d.varType.get).sym,
                d);
        d.sym.type = t;
    }

    override void visitStructDecl(StructDecl s)
    {
        auto old = current.symbol;
        current.symbol = s.sym;
        inFunctionBodyStack.push(false);
        super.visitStructDecl(s);
        inFunctionBodyStack.pop();
        current.symbol = old;
    }

    override void visitClassDecl(ClassDecl s)
    {
        auto old = current.symbol;
        current.symbol = s.sym;
        inFunctionBodyStack.push(false);
        super.visitClassDecl(s);
        inFunctionBodyStack.pop();
        current.symbol = old;
    }

    override void visitInterfaceDecl(InterfaceDecl s)
    {
        auto old = current.symbol;
        current.symbol = s.sym;
        inFunctionBodyStack.push(false);
        super.visitInterfaceDecl(s);
        inFunctionBodyStack.pop();
        current.symbol = old;
    }

    DType typeOf(Identifier id, Scope sc)
    {
        if(auto i = cast(PointerIdentifier)id)
            return (typeOf(i.pointerOf, sc)).getPointerTo();
        else if(auto i = cast(StaticArrayIdentifier)id)
            return typeOf(i.arrayOf, sc).getAsStaticArray(i.size);
        return sc.findType(id.get);
    }

    Module[] modules;
    Module current;
    SmallArray!(bool) inFunctionBodyStack;
}

class TypeBuilder : Visitor!(void)
{
    override void visit(Module[] modules)
    {
        foreach (mod; modules)
        {
            current = mod;
            foreach (decl; mod.decls)
                visitDecl(decl);
        }
    }

    override void visitStructDecl(StructDecl s)
    {
        auto st = s.env.findType(s.identifier.get).asStruct;
        s.sym = current.symbol.createMember(
                s.identifier.get, 
                st,
                s.env.find(s.identifier.get));

        foreach (decl; s.decls)
        {
            DType type;
            char[] name;
            if (auto varDecl = cast(VarDecl)decl)
            {
                type = typeOf(varDecl.varType, varDecl.env);
                name = varDecl.identifier.get;
            }
            else if (auto fd = cast(FuncDecl)decl)
            {
                type = fd.type;
                name = fd.identifier.get;
            }
            st.addMember(type, name);
        }
    }

    override void visitClassDecl(ClassDecl s)
    {
        auto st = s.env.findType(s.identifier.get).asClass;
        s.sym = current.symbol.createMember(
                s.identifier.get, 
                st,
                s.env.find(s.identifier.get));

        foreach (decl; s.decls)
        {
            DType type;
            char[] name;
            if (auto varDecl = cast(VarDecl)decl)
            {
                type = typeOf(varDecl.varType, varDecl.env);
                name = varDecl.identifier.get;
            }
            else if (auto fd = cast(FuncDecl)decl)
            {
                type = fd.type;
                name = fd.identifier.get;
            }
            st.addMember(type, name);
        }
    }

    override void visitInterfaceDecl(InterfaceDecl s)
    {
        auto st = s.env.findType(s.identifier.get).asInterface;
        s.sym = current.symbol.createMember(
                s.identifier.get, 
                st,
                s.env.find(s.identifier.get));

        foreach (decl; s.decls)
        {
            DType type;
            char[] name;
            if (auto varDecl = cast(VarDecl)decl)
            {
                type = typeOf(varDecl.varType, varDecl.env);
                name = varDecl.identifier.get;
            }
            else if (auto fd = cast(FuncDecl)decl)
            {
                type = fd.type;
                name = fd.identifier.get;
            }
            st.addMember(type, name);
        }
    }


    DType typeOf(Identifier id, Scope sc)
    {
        if(auto i = cast(PointerIdentifier)id)
            return (typeOf(i.pointerOf, sc)).getPointerTo();
        else if(auto i = cast(StaticArrayIdentifier)id)
            return typeOf(i.arrayOf, sc).getAsStaticArray(i.size);
        return sc.findType(id.get);
    }

    Module current;
}


/**
  Add scopes to everything, and add all identifiers that correspond to types.
  Types/Symbols are added by ForwardReference.
 **/
class ScopeBuilder : Visitor!(void)
{
    static ModuleHandler mHandle;

    static this()
    {
        mHandle = new ModuleHandler;
    }

    this()
    {
    }

    override void visit(Module[] modules)
    {
        foreach(m ; modules)
            visitModule(m);

        auto fr = new ForwardReference();

        fr.visit(modules);
    }

    private void registerBasicTypeTo(char[] n, DType t, Scope sc, Module m)
    {
        sc.types[n] = t;
        auto id = new Identifier(n);
        id.env = sc;
        auto decl = new DummyDecl();
        auto sym = m.symbol.createMember(n, t, decl);
        sym.decl = decl;
        decl.sym = sym;
        decl.env = sc;
        sc.put(id.get, decl);
    }

    override void visitModule(Module m)
    {
        auto root = new Scope;
        table ~= root;

        m.symbol = new Symbol;

        registerBasicTypeTo("void",     DType.Void, root, m);
        registerBasicTypeTo("bool",     DType.Bool, root, m);
        registerBasicTypeTo("byte",     DType.Byte, root, m);
        registerBasicTypeTo("ubyte",    DType.UByte, root, m);
        registerBasicTypeTo("short",    DType.Short, root, m);
        registerBasicTypeTo("ushort",   DType.UShort, root, m);
        registerBasicTypeTo("int",      DType.Int, root, m);
        registerBasicTypeTo("uint",     DType.UInt, root, m);
        registerBasicTypeTo("long",     DType.Long, root, m);
        registerBasicTypeTo("ulong",    DType.ULong, root, m);

        registerBasicTypeTo("char",     DType.Char, root, m);
        registerBasicTypeTo("wchar",    DType.WChar, root, m);
        registerBasicTypeTo("dchar",    DType.DChar, root, m);

        registerBasicTypeTo("float",    DType.Float, root, m);
        registerBasicTypeTo("double",   DType.Double, root, m);
        registerBasicTypeTo("real",     DType.Real, root, m);

        current().inModule = m;
        current().mHandle = mHandle;
        mHandle.add(m);
        m.env = current();
        super.visitModule(m);
    }

    override void visitDecl(Decl d)
    {
        d.env = current();
        super.visitDecl(d);
    }

    override void visitImportDecl(ImportDecl i)
    {
        i.env.imports ~= i;
        super.visitImportDecl(i);
    }

    override void visitStmt(Stmt s)
    {
        s.env = current();
        s.stmtIndex = s.env.stmtIndex;
        super.visitStmt(s);
    }

    override void visitExp(Exp e)
    {
        e.env = current();
        e.stmtIndex = e.env.stmtIndex;
        super.visitExp(e);
    }

    override void visitFuncDecl(FuncDecl d)
    {
        current().put(d.identifier.get, d);
        d.env = current();
        auto sc = push();

        visitExp(d.returnType);
        visitExp(d.identifier);
        sc.parentFunction = d;
        foreach (arg; d.funcArgs)
            visitDecl(arg);
        foreach (stmt; d.statements)
        {
            sc.currentStmtIndex++;
            visitStmt(stmt);
        }
        pop(sc);
    }

    override void visitVarDecl(VarDecl d)
    {
        if (d.init)
            visitExp(d.init);

        if (need_push > 0 && current().parentFunction !is null) {
            push();
            --need_push;
        }

        auto sc = current();
        sc.put(d.identifier.get, d);
        d.env = sc;
        visitExp(d.varType);
        visitExp(d.identifier);
    }

    override void visitStructDecl(StructDecl s)
    {
        auto sc = current();
        sc.put(s.identifier.get, s);
        s.env = sc;
        auto type = new DStruct(s.identifier);

        sc.types[s.identifier.get] = type;

        sc = push();
        super.visitStructDecl(s);
        pop(sc);
    }

    override void visitClassDecl(ClassDecl s)
    {
        auto sc = current();
        sc.put(s.identifier.get, s);
        s.env = sc;
        auto type = new DClass(s.identifier);

        sc.types[s.identifier.get] = type;

        sc = push();
        super.visitClassDecl(s);
        pop(sc);
    }

    override void visitInterfaceDecl(InterfaceDecl s)
    {
        auto sc = current();
        sc.put(s.identifier.get, s);
        s.env = sc;
        auto type = new DInterface(s.identifier);

        sc.types[s.identifier.get] = type;

        sc = push();
        super.visitInterfaceDecl(s);
        pop(sc);
    }

    override void visitDeclStmt(DeclStmt d)
    {
        ++need_push;
        super.visitDeclStmt(d);
    }
    private uint need_push = 0;

    override void visitIfStmt(IfStmt s)
    {
        s.env = current();
        visitExp(s.cond);
        auto sc = push();
        visitStmt(s.then_body);
        pop(sc);

        if (s.else_body !is null)
        {
            sc = push();
            visitStmt(s.else_body);
            pop(sc);
        }
    }

    override void visitWhileStmt(WhileStmt s)
    {
        s.env = current();
        auto sc = push();
        super.visitWhileStmt(s);
        pop(sc);
    }

    override void visitForStmt(ForStmt s)
    {
        s.env = current();
        auto sc = push();
        super.visitForStmt(s);
        pop(sc);
    }

    override void visitCompoundStmt(CompoundStatement s)
    {
        s.env = current();
        auto sc = push();
        super.visitCompoundStmt(s);
        pop(sc);
    }

private:
    Scope[] table;

    Scope push()
    {
        auto sc = new Scope(current());
        table ~= sc;
        return sc;
    }

    Scope pop(Scope sc = null)
    {
        if (sc !is null)
        {
            table.length = table.find(sc);
            return sc;
        }

        auto res = table[$ - 1];
        table.length = table.length - 1;
        return res;
    }

    Scope current()
    {
        return table[$ - 1];
    }
}