// ir-restructure.cpp
#include "slang-ir-restructure.h"

#include "slang-ir-insts.h"
#include "slang-ir.h"

namespace Slang
{
bool Region::isDescendentOf(Region* other)
{
    Region* rr = this;
    while (rr)
    {
        if (rr == other)
            return true;

        rr = rr->getParent();
    }
    return false;
}

bool Region::isDescendentOf(IRBlock* block)
{
    Region* rr = this;
    while (rr)
    {
        if (rr->getFlavor() == Region::Flavor::Simple)
        {
            SimpleRegion* simpleRegion = (SimpleRegion*)rr;
            if (simpleRegion->block == block)
                return true;
        }

        rr = rr->getParent();
    }
    return false;
}

/// An "active" label during control flow (re)structuring.
struct LabelStack
{
    /// Possible operations associated with labels.
    enum class Op
    {
        Break,
        Continue,

        CountOf,
    };

    /// What kind of operation does a branch to this label represent?
    Op op;

    /// The next label down on the stack
    LabelStack* parent;

    /// The block the represents this label in the IR control flow graph.
    IRBlock* block;

    /// The region that represents this label in the structured program
    Region* region;
};

/// State used when restructuring control flow.
struct ControlFlowRestructuringContext
{
    /// Sink to use when diagnosing errors in control-flow restructuring.
    ///
    /// The restructuring pass should be able to handle anything the front-end
    /// throws at it, so these errors will all be unexpected. Still, we need
    /// a way to report them cleanly without crashing the process.
    ///
    DiagnosticSink* sink = nullptr;
    DiagnosticSink* getSink() { return sink; }

    /// The region tree we are in the process of building.
    RegionTree* regionTree = nullptr;

    /// Whether to preserve fall-through structure in switch statements.
    /// When false, fall-through is not preserved (legacy behavior for HLSL/WGSL).
    bool preserveFallThrough = true;
};

/// Check if a block or any of its reachable blocks (within a case)
/// branches unconditionally to the target block.
/// This is used to detect if a switch case falls through to the next case.
/// Uses iterative traversal with explicit worklist to avoid stack overflow.
static bool caseBlocksFallThroughTo(
    IRBlock* startBlock,
    IRBlock* targetBlock,
    IRBlock* breakLabel,
    HashSet<IRBlock*>& visited)
{
    if (!startBlock || !targetBlock)
        return false;

    List<IRBlock*> worklist;
    worklist.add(startBlock);

    while (worklist.getCount() > 0)
    {
        auto block = worklist.getLast();
        worklist.removeLast();

        if (!block)
            continue;

        if (visited.contains(block))
            continue;
        visited.add(block);

        // Don't follow branches to the break label
        if (block == breakLabel)
            continue;

        // If we reached the target, there's fall-through
        if (block == targetBlock)
            return true;

        auto terminator = block->getTerminator();
        if (!terminator)
            continue;

        switch (terminator->getOp())
        {
        case kIROp_UnconditionalBranch:
            {
                auto branch = as<IRUnconditionalBranch>(terminator);
                auto target = branch->getTargetBlock();
                if (target == targetBlock)
                    return true;
                if (target != breakLabel)
                    worklist.add(target);
            }
            break;
        case kIROp_ConditionalBranch:
        case kIROp_IfElse:
            {
                auto branch = as<IRConditionalBranch>(terminator);
                worklist.add(branch->getTrueBlock());
                worklist.add(branch->getFalseBlock());
            }
            break;
        case kIROp_Loop:
            {
                // For loops, follow the break block (where control goes after the loop)
                auto loop = as<IRLoop>(terminator);
                worklist.add(loop->getBreakBlock());
            }
            break;
        case kIROp_Switch:
            {
                // For nested switches, follow the break label (where control goes after)
                auto nestedSwitch = as<IRSwitch>(terminator);
                worklist.add(nestedSwitch->getBreakLabel());
            }
            break;
        default:
            // Return, throw, etc. - don't follow
            break;
        }
    }

    return false;
}

/// Convert a range of blocks in the IR CFG into a region.
///
/// We want to generate a region that stands in for the
/// blocks that are logically in the internal [begin, end)
/// which we consider as representing a single-entry multiple-exit
/// sub-graph. Note that `end` is *not* part of the sub-graph,
/// but instead points to a block that is logically "after"
/// the sub-graph. `end` can be `null` to indicate that the
/// sub-graph extends as far as possible.
///
/// Because there can be multiple exits, control flow may
/// exit the sub-graph without branching to `end`, any
/// such "non-local" branching should be to one of the
/// blocks stored in the current `LabelStack`.
///
// TODO: Eventually we should replace all of this logic with
// a variation on the "Relooper" algorithm as it is used
// in Emscripten.
//
static RefPtr<Region> generateRegionsForIRBlocks(
    ControlFlowRestructuringContext* ctx,
    Region* inParentRegion,
    IRBlock* begin,
    IRBlock* end,
    LabelStack* initialLabels,    // Labels to use at the start
    LabelStack* labels = nullptr) // Labels to switch to after emitting first basic block
{
    if (!labels)
        labels = initialLabels;
    auto useLabels = initialLabels;

    //
    // We will try to build up as long of a sequential/simple region
    // as possible, to avoid deep recursion in this algorithm.
    //
    RefPtr<Region> resultRegion = nullptr;
    RefPtr<Region>* resultLink = &resultRegion;

    // As we move along, the parent region to use for regions
    // we create will shift, so we need a temporary to track
    // the current parent region.
    //
    Region* parentRegion = inParentRegion;

    //
    // We will start with the `begin` block, and try to proceed
    // sequentially until we see the `end` block, or run into
    // an edge that exits teh region.
    //
    IRBlock* block = begin;
    while (block != end)
    {
        // If the block we are trying to emit has been registered as a
        // destination label (e.g. for a loop or `switch`) then we
        // need to exit the current region, which amounts to generating
        // a `break` or `continue` operation.
        //
        // TODO: we eventually need to handle the possibility of
        // multi-level break/continue targets, which could be challenging.

        // Because we will only support single-level break/continue, we
        // want to resolve what is the most recent label that is "active"
        // for the given operation (`break` or `continue`).
        //
        // We will do this with a naive loop, just to keep things simple.
        // We start with no block "regsitered" as the target for each
        // operation.
        //
        IRBlock* registeredBlock[(int)LabelStack::Op::CountOf] = {};
        for (auto ll = useLabels; ll; ll = ll->parent)
        {
            // For each active label, see if it is the first one
            // we encounter for the given op.
            //
            if (!registeredBlock[(int)ll->op])
            {
                registeredBlock[(int)ll->op] = ll->block;
            }
        }

        // Next we will search through *all* of the registered labels,
        // and see if one of them matches the current `block`.
        //
        for (auto ll = useLabels; ll; ll = ll->parent)
        {
            // Does this label match the block we are trying to translate?
            if (ll->block != block)
                continue;

            // Okay, the block we are trying to generate code for is a label
            // that we should branch to (we shouldn't just emit the code here
            // and now...)
            //
            // We should first confirm that the block is the inner-most label
            // registered for the given control-flow op (`break` or `continue`)
            // because if it *isn't* we currently can't generate code.
            //
            if (block != registeredBlock[(int)ll->op])
            {
                if (ctx->getSink())
                    ctx->getSink()->diagnose(block, Diagnostics::multiLevelBreakUnsupported);
            }

            // Now we need to create a structured `break` or `continue` operation
            // to match the operation associated with the target.
            //
            switch (ll->op)
            {
            case LabelStack::Op::Break:
                {
                    auto outerRegion = (BreakableRegion*)ll->region;
                    RefPtr<BreakRegion> breakRegion = new BreakRegion(parentRegion, outerRegion);

                    *resultLink = breakRegion;
                    resultLink = nullptr;
                }
                break;

            case LabelStack::Op::Continue:
                {
                    auto outerRegion = (LoopRegion*)ll->region;
                    RefPtr<ContinueRegion> continueRegion =
                        new ContinueRegion(parentRegion, outerRegion);

                    *resultLink = continueRegion;
                    resultLink = nullptr;
                }
                break;
            }

            // If the `block` matched an active label, then we should have
            // created a branch, and there is nothing to be done here.
            return resultRegion;
        }

        // We now know that the given `block` is part of our control-flow region,
        // so we need to output a simple region that executes the code in that block.
        //
        RefPtr<SimpleRegion> simpleRegion = new SimpleRegion(parentRegion, block);

        // We need to register the mapping from `block` to this region, but in
        // general this isn't a one-to-one mapping, but rather one-to-many.
        // This is because a "continue clause" in a `for` loop might get duplicated
        // at each `continue` site in the output code. To deal with this
        // we build a singly-linked list of regions for each block.
        //
        // TODO: confirm that continue clauses are the only case that leads
        // to duplication.
        //
        // TODO: remove this workaround once we have a more powerful restructuring
        // pass that avoids duplicating blocks (by introducing new temporaries...)
        //
        SimpleRegion* nextSimpleRegionForSameBlock = nullptr;
        ctx->regionTree->mapBlockToRegion.tryGetValue(block, nextSimpleRegionForSameBlock);
        ctx->regionTree->mapBlockToRegion[block] = simpleRegion;

        *resultLink = simpleRegion;
        resultLink = &simpleRegion->nextRegion;
        parentRegion = simpleRegion;

        // The simple region we created will represent all of the non-terminator
        // instructions in the `block`, so now we need to figure out what to
        // create to represent that terminator.
        //
        auto terminator = block->getTerminator();
        SLANG_ASSERT(terminator != nullptr);
        switch (terminator->getOp())
        {
        default:
        case kIROp_ConditionalBranch:
            // Note: we don't currently generate ordinary `conditionalBranch` instructions,
            // and instead only generate `ifElse` instructions, which include additional
            // information that can inform our control-flow restructuring pass.
            //
            SLANG_UNEXPECTED("unhandled terminator instruction opcode");
            [[fallthrough]];
        case kIROp_Unreachable:
        case kIROp_MissingReturn:
        case kIROp_Return:
        case kIROp_GenericAsm:
            // These cases are all simple terminators that can be handled as-is
            // without needing to construct a separate `Region` to encapsulate them.
            //
            // We will cap off the current sequence of simple regions and return.
            //
            *resultLink = nullptr;
            return resultRegion;

        case kIROp_IfElse:
            {
                // Here we have a two-way branch, so that we will construct a
                // region representing an `if` statement.
                //
                auto ifInst = (IRIfElse*)terminator;
                auto trueBlock = ifInst->getTrueBlock();
                auto falseBlock = ifInst->getFalseBlock();
                auto afterBlock = ifInst->getAfterBlock();


                RefPtr<IfRegion> ifRegion = new IfRegion(parentRegion, ifInst);

                // The region for the "then" part of things will consist of
                // the range of blocks `[trueBlock, afterBlock)`.
                //
                // This logic assumes that `afterBlock` is a valid structured
                // "join point" such that any branch out of the sub-region
                // either leads to `afterBlock` *or* one of the labels
                // that is already present on our label stack.
                //
                ifRegion->thenRegion =
                    generateRegionsForIRBlocks(ctx, ifRegion, trueBlock, afterBlock, labels);

                // Generating a region for the `else` part is similar.
                // Note that it is possible for this to be a `null`
                // region, if `falseBlock == afterBlock`.
                //
                ifRegion->elseRegion =
                    generateRegionsForIRBlocks(ctx, ifRegion, falseBlock, afterBlock, labels);

                *resultLink = ifRegion;
                resultLink = &ifRegion->nextRegion;
                parentRegion = ifRegion;

                // Continue with the block after the `ifElse` instruction.
                block = afterBlock;
            }
            break;

        case kIROp_Loop:
            {
                // The terminator in this case is the header for a structured loop.
                //
                auto loopInst = (IRLoop*)terminator;
                auto bodyBlock = loopInst->getTargetBlock();
                auto afterBlock = loopInst->getBreakBlock();

                RefPtr<LoopRegion> loopRegion = new LoopRegion(parentRegion, loopInst);

                // We will need to set up entries on our label stack to
                // represent the targets for `break` or `continue`
                // operations inside the loop.
                //
                // First we set up the stack entry for the `break` label,
                // which will refer to the block *after* the loop.
                //
                // The region we specify for the label will still be
                // the loop region, though, because the loop is what
                // we are breaking out of.
                //
                LabelStack loopBreakLabelStack;
                loopBreakLabelStack.parent = labels;
                loopBreakLabelStack.block = afterBlock;
                loopBreakLabelStack.region = loopRegion;
                loopBreakLabelStack.op = LabelStack::Op::Break;

                //
                // The `continue` label warrants a bit more careful explanation,
                // because it will *not* refer to the block that was regsitered
                // as the continue target in the IR `loop` instruction. This
                // is because we will always emit our loops as `for(;;) { ... }`
                // with no continue clause at all, so that a `continue` in
                // the output code will always refer to the top of the loop.
                //
                // This means that the `continue` label for the purposes of
                // structured control flow will be the start of the loop body:
                //
                LabelStack loopContinueLabelStack;
                loopContinueLabelStack.parent = &loopBreakLabelStack;
                loopContinueLabelStack.block = bodyBlock;
                loopContinueLabelStack.region = loopRegion;
                loopContinueLabelStack.op = LabelStack::Op::Continue;
                //
                // Note: by ignoring the original continue block from the
                // high-level loop, we create a situation where that code
                // might get emitted more than once (once per implicit
                // or explicit `continue` site in the original program).
                //
                // That is an acceptable trade-off for now, because continue
                // blocks will usually be small (and fxc makes the same choice),
                // but it could lead to Bad Things if somebody were to call
                // a function in their continue clause, and that function does
                // a compute shader barrier operation.
                //
                // A better long-term fix is to take a high-level loop like:
                //
                //      for(A; B; C) { ... continue; ... break; ... }
                //
                // and translate it into something like the following (assuming
                // we have labeled statements and multi-level `break`):
                //
                //      A;
                //      Outer: for(;;) {
                //          Inner: for(;;) {
                //              if(B) {} else break Outer;
                //              ...
                //              break Inner; // `continue` becomes break of inner loop
                //              ...
                //              break Outer; // `break` becomes break of outer loop
                //              ...
                //              break; // inner loop unconditionally breaks at the end
                //          }
                //          C;  // continue clause comes after inner loop
                //      }
                //
                // If you draw up a control flow graph for that code, you'll find
                // it is equivalent to the orignal `for` loop, but now supports
                // arbitrary code (not just a single expression) for the continue clause.
                // Unlike the current code-duplication solution, `C` appears only once
                // in the output, and seems to clearly be at a "joint point" for control
                // flow so that it is clear that a barrier there is valid in GLSL.
                //
                // Anyway, back our regularly scheduled programming.
                //
                // With the label stack stuff set up, we want to take the region
                // of the CFG defined by `[bodyBlock, afterBlock)` and turn it into
                // the body region for our loop.
                //
                // The only thing we want to be a little bit careful about is
                // that we don't want the logic at the top of this function
                // that looks for a block it can translate into a `continue`
                // to trigger on `bodyBlock`, since that means we'd just turn
                // the whole body into a single `continue`.
                //
                // To avoid this problem, we pass in two different label stacks:
                // one to use for the first block, and one to use for subsequent
                // blocks.
                //
                loopRegion->body = generateRegionsForIRBlocks(
                    ctx,
                    loopRegion,
                    bodyBlock,
                    // TODO: should we pass `afterBlock` here instead of `null`?
                    nullptr,
                    // For the first block, we only want the `break` label active
                    &loopBreakLabelStack,
                    // After the first block, we can safely use the `continue` label too
                    &loopContinueLabelStack);

                *resultLink = loopRegion;
                resultLink = &loopRegion->nextRegion;
                parentRegion = loopRegion;

                // Continue with the block after the loop
                block = afterBlock;
            }
            break;

        case kIROp_UnconditionalBranch:
            {
                // Here we have an unconditional branch that was
                // not covered by one of our labels for non-local
                // branches (`break` or `continue`).
                //
                // We will thus assume that the target of the
                // branch is part of the same region we are building,
                // and continue with the target block;
                //
                auto branchInst = (IRUnconditionalBranch*)terminator;
                block = branchInst->getTargetBlock();
            }
            break;

        case kIROp_Switch:
            {
                // A `switch` instruction will always translate
                // to a `SwitchRegion` and then to a `switch` statement.
                //
                // We will need to take care to emit `case`s in ways
                // that avoid code duplication.
                //
                // The logic here isn't going to be robust in edge cases
                // (please don't write Duff's Device in Slang just yet).
                // Doing significantly better than what is here would
                // require something like the Relooper algorithm, though.
                //
                auto switchInst = (IRSwitch*)terminator;
                auto breakLabel = switchInst->getBreakLabel();
                auto defaultLabel = switchInst->getDefaultLabel();

                RefPtr<SwitchRegion> switchRegion = new SwitchRegion(parentRegion, switchInst);

                // A direct branch to the block after the `switch` can
                // be emitted as a `break` statement, so we will register
                // the appropriate label on a label stack:
                //
                LabelStack switchBreakLabelStack;
                switchBreakLabelStack.parent = labels;
                switchBreakLabelStack.op = LabelStack::Op::Break;
                switchBreakLabelStack.block = breakLabel;
                switchBreakLabelStack.region = switchRegion;

                // Track whether we've warned about fall-through being restructured
                bool warnedAboutFallthrough = false;

                // Build a source-ordered list of switch case blocks.
                // This is necessary to correctly handle fall-through when default
                // appears before explicit cases in source order (e.g., default -> case 0).
                //
                // We collect all unique case/default blocks with their associated values,
                // then sort by their position in the function to get source order.

                struct SwitchCaseInfo
                {
                    IRBlock* label;
                    List<IRInst*> values; // empty for default-only
                    bool isDefault;
                };

                // Build a map from block to position in function (source order)
                Dictionary<IRBlock*, Index> blockOrder;
                {
                    Index orderIdx = 0;
                    auto func = cast<IRGlobalValueWithCode>(block->getParent());
                    for (auto b = func->getFirstBlock(); b; b = b->getNextBlock())
                    {
                        blockOrder[b] = orderIdx++;
                    }
                }

                // Collect all unique case blocks with their values
                Dictionary<IRBlock*, Index> blockToInfoIndex;
                List<SwitchCaseInfo> caseInfos;

                UInt caseCount = switchInst->getCaseCount();
                for (UInt i = 0; i < caseCount; i++)
                {
                    auto label = switchInst->getCaseLabel(i);
                    auto value = switchInst->getCaseValue(i);

                    Index infoIdx;
                    if (blockToInfoIndex.tryGetValue(label, infoIdx))
                    {
                        // Multiple case values share the same label
                        caseInfos[infoIdx].values.add(value);
                    }
                    else
                    {
                        SwitchCaseInfo info;
                        info.label = label;
                        info.values.add(value);
                        info.isDefault = (label == defaultLabel);
                        caseInfos.add(info);
                        blockToInfoIndex[label] = caseInfos.getCount() - 1;
                    }
                }

                // If default wasn't already added (not sharing label with a case)
                // and default actually has a body (doesn't just branch to break)
                if (defaultLabel != breakLabel && !blockToInfoIndex.containsKey(defaultLabel))
                {
                    SwitchCaseInfo info;
                    info.label = defaultLabel;
                    info.isDefault = true;
                    caseInfos.add(info);
                }

                // Sort by block position in function (source order)
                caseInfos.sort([&](const SwitchCaseInfo& a, const SwitchCaseInfo& b)
                               { return blockOrder[a.label] < blockOrder[b.label]; });

                // Now iterate through cases in source order
                for (Index caseIdx = 0; caseIdx < caseInfos.getCount(); caseIdx++)
                {
                    auto& info = caseInfos[caseIdx];

                    RefPtr<SwitchRegion::Case> currentCase = new SwitchRegion::Case();
                    switchRegion->cases.add(currentCase);

                    // Add case values
                    for (auto val : info.values)
                        currentCase->values.add(val);

                    // Track if this is the default
                    if (info.isDefault)
                    {
                        switchRegion->defaultCase = currentCase;
                    }

                    // Now we need to generate a region for the instructions
                    // that make up this case. The 99% case will be that it
                    // will terminate with a `break` (or a `return`,
                    // `continue`, etc.) and so we can pass in `nullptr`
                    // for the ending block.
                    //
                    IRBlock* caseEndLabel = nullptr;

                    // However, there is also the possibility that
                    // this `case` will fall through to the next, and
                    // so we need to prepare for that possibility here.
                    //
                    // If there *is* a next case/default in source order, then we will
                    // set its label up as the "end" label when emitting the statements
                    // inside the block.
                    //
                    // However, if the target doesn't support fall-through (e.g., HLSL/WGSL),
                    // we should NOT set the end label to the next case, so that each case
                    // body will be treated as terminating independently.
                    IRBlock* potentialFallThroughTarget = nullptr;
                    if (caseIdx + 1 < caseInfos.getCount())
                    {
                        potentialFallThroughTarget = caseInfos[caseIdx + 1].label;
                    }

                    if (ctx->preserveFallThrough && potentialFallThroughTarget)
                    {
                        caseEndLabel = potentialFallThroughTarget;
                    }
                    else if (
                        !ctx->preserveFallThrough && potentialFallThroughTarget &&
                        !warnedAboutFallthrough && ctx->getSink())
                    {
                        // Check if this case actually falls through
                        HashSet<IRBlock*> visited;
                        if (caseBlocksFallThroughTo(
                                info.label,
                                potentialFallThroughTarget,
                                breakLabel,
                                visited))
                        {
                            ctx->getSink()->diagnose(
                                switchInst,
                                Diagnostics::switchFallthroughRestructured);
                            warnedAboutFallthrough = true;
                        }
                    }

                    // Now we can actually generate the region.
                    //
                    currentCase->body = generateRegionsForIRBlocks(
                        ctx,
                        switchRegion,
                        info.label,
                        caseEndLabel,
                        &switchBreakLabelStack);
                }

                *resultLink = switchRegion;
                resultLink = &switchRegion->nextRegion;
                parentRegion = switchRegion;

                // Continue with the block after the `switch`
                block = breakLabel;
            }
            break;
        }

        // After we've emitted the first block, we are safe from accidental
        // cases where we'd emit an entire loop body as a single `continue`,
        // so we can safely switch in whatever labels are intended to be used.
        useLabels = labels;

        // If we reach this point, then we've emitted
        // one block, and we have a new block where
        // control flow continues.
        //
        // We need to handle a special case here,
        // when control flow jumps back to the
        // starting block of the range we were
        // asked to work with:
        if (block == begin)
        {
            break;
        }
    }

    // We seem to have reached the rend of the region
    // without anything special happening. This means
    // we should cap off the current sequence of regions
    // and return what we have.
    //
    *resultLink = nullptr;
    return resultRegion;
}

RefPtr<RegionTree> generateRegionTreeForFunc(
    IRGlobalValueWithCode* code,
    DiagnosticSink* sink,
    bool preserveFallThrough)
{
    RefPtr<RegionTree> regionTree = new RegionTree();
    regionTree->irCode = code;

    ControlFlowRestructuringContext restructuringContext;
    restructuringContext.sink = sink;
    restructuringContext.regionTree = regionTree;
    restructuringContext.preserveFallThrough = preserveFallThrough;

    regionTree->rootRegion = generateRegionsForIRBlocks(
        &restructuringContext,
        nullptr,
        code->getFirstBlock(),
        nullptr,
        nullptr);

    return regionTree;
}
} // namespace Slang
