package com.seaofnodes.simple.node;

import com.seaofnodes.simple.codegen.CodeGen;
import com.seaofnodes.simple.type.*;
import com.seaofnodes.simple.util.Ary;

import java.util.BitSet;
import java.util.IdentityHashMap;

/** Control Flow Graph Nodes
 * <p>
 *  CFG nodes have a immediate dominator depth (idepth) and a loop nesting
 *  depth(loop_depth).
 * <p>
 *  idepth is computed lazily upon first request, and is valid even in the
 *  Parser, and is used by peepholes during parsing and afterward.
 * <p>
 *  loop_depth is computed after optimization as part of scheduling.
 *
 *  Start - Block head; has constants and trailing Funs
 *  Region - Block head; has Phis.  Includes Fun and Parms.
 *  CallEnd - Block head; followed by CProj and Proj; also linked RPC Parms
 *  CProj - Block head
 *
 *  If - Block tail; followed by CProj same block
 *  Call - Block tail; followed by CallEnd and linked Funs
 *  New/intrinsic - If followed by CProj: block tail else mid-block; followed by Proj
 *
 */
public abstract class CFGNode extends Node {

    public CFGNode(Node...   nodes) { super(nodes); }
    public CFGNode(CFGNode cfg) {
        super(cfg);
        if( cfg != null ) {
            _idepth = cfg._idepth;
            _idepthVersion = cfg._idepthVersion;
            _ltree = cfg._ltree;
            _pre = cfg._pre;
        }
    }

    public CFGNode cfg(int idx) { return (CFGNode)in(idx); }

    // Block head is Start, Region, CProj, but not e.g. If, Return, Stop
    public boolean blockHead() { return false; }

    // Get the one control following; error to call with more than one e.g. an
    // IfNode or other multi-way branch.
    public CFGNode uctrl() {
        CFGNode c = null;
        for( Node n : _outputs )
            if( n instanceof CFGNode cfg )
                {  assert c==null;  c = cfg; }
        return c;
    }

    // Used by the encoding / final BB layout
    public CFGNode uctrlSkipEmpty() {
        CFGNode x = this, y;
        while( x.nOuts() == 1 && (y=x.uctrl())!=null && !(y instanceof CallNode) ) // Skip empty blocks
            x = y;
        return x;
    }

    // ------------------------------------------------------------------------
    /**
     * Immediate dominator tree depth, used to approximate a real IDOM during
     * parsing where we do not have the whole program, and also peepholes
     * change the CFG incrementally.
     * <p>
     * See {@link <a href="https://en.wikipedia.org/wiki/Dominator_(graph_theory)">...</a>}
     */
    public char _idepth;
    public char _idepthVersion;
    public int idepth() { return validIDepth() ? _idepth : cacheIDepth(idom().idepth()+1); }
    // Zero depth means uncached. Check before narrowing so overflow cannot wrap.
    final int cacheIDepth(int depth) {
        assert 0 <= depth && depth <= Character.MAX_VALUE : "Dominator depth exceeds 65535";
        _idepthVersion = CodeGen.CODE.iDepthVersion();
        return _idepth = (char)depth;
    }

    final boolean validIDepth() { return _idepth!=0 && _idepthVersion==CodeGen.CODE.iDepthVersion(); }


    // Return the immediate dominator of this Node and compute dom tree depth.
    public CFGNode idom(Node dep) { return cfg(0); }
    public final CFGNode idom() { return idom(null); }

    // Return the LCA of two idoms
    public CFGNode domLCA(CFGNode rhs, Node dep) {
        if( rhs==null ) return this;
        CFGNode lhs = this;
        while( lhs != rhs ) {
            if( lhs==null || rhs==null ) return null;
            int comp = lhs.idepth() - rhs.idepth();
            if( comp >= 0 ) lhs = (dep==null ? lhs : dep.addDep(lhs)).idom();
            if( comp <= 0 ) rhs = (dep==null ? rhs : dep.addDep(rhs)).idom();
        }
        return lhs;
    }

    // Safe read of the control input of a data node.
    public static CFGNode safeCFG(Node n) {
        if( n==null ) return null;
        if( !(n.in(0) instanceof CFGNode cfg) ) return null;
        return cfg._safeCFG();
    }
    private CFGNode _safeCFG() {
        return switch(this) {
        case CallNode call when call.cend()._folding -> call.cfg0()._safeCFG();
        case CallEndNode cend when cend._folding -> cend.cfg(1)._safeCFG();
        case FunNode fun when fun._folding -> fun.cfg(1)._safeCFG();
        case ReturnNode ret when ret.fun()._folding -> ret.cfg0()._safeCFG();
        default -> this;
        };
    }

    // True when this node dominates `sub` in the current idom tree.
    public boolean dominates(CFGNode sub, Node dep) {
        for( CFGNode cfg=sub; cfg!=null; cfg=cfg.idom() )
            if( cfg == this )
                return true;
            else if( dep != null ) dep.addDep(cfg);
        return false;
    }

    public static CFGNode deepest(CFGNode a, CFGNode b) {
        if( a == null ) return b;
        if( b == null ) return a;
        return b.idepth() > a.idepth() ? b : a;
    }

    public static CFGNode earlyCFG(Node n) {
        return earlyCFG(n,new IdentityHashMap<>(),new BitSet());
    }

    public static CFGNode earlyCFG(Node n, IdentityHashMap<Node,CFGNode> cache, BitSet active) {
        if( n == null ) return null;
        CFGNode cached = cache.get(n);
        if( cached != null ) return cached;
        if( active.get(n._nid) )
            return safeCFG(n);
        active.set(n._nid);
        CFGNode early = null;
        for( Node x : n._inputs )
            early = deepest(early,defCFG(x,cache,active));
        active.clear(n._nid);
        cache.put(n,early);
        return early;
    }

    private static CFGNode defCFG(Node n, IdentityHashMap<Node,CFGNode> cache, BitSet active) {
        if( n == null || n.isConst() )  return null; // no impact on early cfg
        return switch( n ) {
        case PhiNode phi -> phi.region();
        case CFGNode cfg -> cfg;
        case ProjNode projNode -> safeCFG( n );
        default -> earlyCFG( n, cache, active );
        };
    }

    public boolean sameFun(CFGNode that) {
        FunNode fun = fun();
        return fun != null && fun == that.fun();
    }

    // Find nearest enclosing FunNode
    public FunNode fun() {
        CFGNode cfg = this;
        while( !(cfg instanceof FunNode fun) ) {
            if( cfg==null ) return null;
            cfg = cfg.idom();
        }
        return fun;
    }

    // ------------------------------------------------------------------------
    // Loop nesting
    public LoopNode loop() { return _ltree._head; }
    public int loopDepth() { return _ltree==null ? 0 : _ltree.depth(); }

    public LoopTree _ltree;
    public int _pre;            // Pre-order numbers for loop tree finding
    public static class LoopTree {
        public LoopTree _par;
        public final LoopNode _head;
        public int _depth;
        LoopTree(LoopNode head) { _head = head; }
        @Override public String toString() { return label(_head); }
        public int depth() {
            return _depth==0 ? (_par==null ? 0 : (_depth = _par.depth()+1)) : _depth;
        }
    }

    // ------------------------------------------------------------------------
    // Tag all CFG Nodes with their containing LoopNode; LoopNodes themselves
    // also refer to *their* containing LoopNode, as well as have their depth.
    // Start contains the function bodies at depth 0; real loops add depth.
    public void buildLoopTree( Ary<FunNode> funs, StopNode stop) {
        _ltree = stop._ltree = CodeGen.CODE.XCTRL._ltree = new LoopTree((StartNode)this);
        Ary<FunNode> work = new Ary<>(FunNode.class);

        BitSet post = new BitSet();
        int pre = 1;
        _pre = pre++;
        for( FunNode fun : funs ) {
            if( fun!=null && !fun.isDead() && !fun._type.isHigh() ) {
                work.push(fun);
                if(fun.in(1) instanceof StartCUNode start && start._pre==0 ) {
                    start._pre = pre++;
                    start._ltree = start.cfg(1)._ltree = _ltree;
                    post.set(start._nid);
                }
            };
        }

        for( int i=0; i<work._len; i++ )
            pre = work.at(i)._bltWalk(pre,work.at(i),stop, post, work);
    }

    int _bltWalk( int pre, FunNode fun, StopNode stop, BitSet post, Ary<FunNode> work ) {
        // Pre-walked?
        if( _pre!=0 ) return pre;
        _pre = pre++;
        // Seed the function's outer tree before walking its body.  Synthetic
        // exits can be added before the post-order walk reaches the entry.
        if( this instanceof FunNode ) _ltree = stop._ltree;
        if( this instanceof ReturnNode ) {
            _ltree = fun._ltree;
            post.set(_nid);
            return pre; // Stop/Start links describe the outside world, not a loop.
        }
        // Pre-walk
        for( Node use : _outputs )
            if( use instanceof CFGNode usecfg ) {
                if( skip( usecfg ) ) {
                    if( usecfg._pre==0 && usecfg instanceof FunNode fun2 ) work.add( fun2 );
                } else
                    pre = usecfg._bltWalk( pre, use instanceof FunNode fuse ? fuse : fun, stop, post, work );
            }

        // Post-order work: find innermost loop
        LoopTree inner = null, ltree;
        for( Node use : _outputs ) {
            if( !(use instanceof CFGNode usecfg) ) continue;
            if( skip(usecfg) ) continue;
            if( usecfg._type == Type.XCONTROL ||       // Do not walk dead control
                usecfg._type == TypeTuple.IF_NEITHER ) // Nor dead IFs
                continue;
            // Child visited but not post-visited?
            if( !post.get(usecfg._nid) ) {
                // Must be a backedge to a LoopNode then
                ltree = usecfg._ltree = new LoopTree((LoopNode)usecfg);
            } else {
                // Take child's loop choice, which must exist
                ltree = usecfg._ltree;
                // If falling into a loop, use the target loop's parent instead
                if( ltree._head == usecfg ) {
                    if( ltree._par == null )
                        // This loop never had an If test choose to take its
                        // exit, i.e. it is a no-exit infinite loop.
                        ltree._par = ltree._head.forceExit(fun,stop)._ltree;
                    ltree = ltree._par;
                }
            }
            // Sort inner loops.  The decision point is some branch far removed
            // from either loop head OR either backedge so requires pre-order
            // numbers to figure out innermost.
            if( inner == null ) { inner = ltree; continue; }
            if( inner == ltree ) continue; // No change
            LoopTree outer = ltree._head._pre > inner._head._pre ? inner : ltree;
            inner =          ltree._head._pre > inner._head._pre ? ltree : inner;
            inner._par = outer;
        }
        // Set selected loop
        if( inner!=null )
            _ltree = inner;
        // Tag as post-walked
        post.set(_nid);
        return pre;
    }

    public boolean skip(CFGNode usecfg) {
        // Only walk control users that are alive.
        // Do not walk from a Call to linked Fun's.
        return usecfg instanceof XCtrlNode ||
                (this instanceof CallNode && usecfg instanceof FunNode) ||
                (this instanceof ReturnNode && usecfg instanceof CallEndNode);
    }

    public static String label( CFGNode target ) {
        return switch(target) {
        case StartNode start -> start.label();
        case FunNode fun -> fun.label();
        case LoopNode loop -> "LOOP"+target._nid;
        default -> "L"+target._nid;
        };
    }

}