// Copyright 2020 The Tint Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef SRC_TINT_READER_SPIRV_CONSTRUCT_H_
#define SRC_TINT_READER_SPIRV_CONSTRUCT_H_

#include <memory>
#include <sstream>
#include <string>
#include <vector>

namespace tint::reader::spirv {

/// A structured control flow construct, consisting of a set of basic blocks.
/// A construct is a span of blocks in the computed block order,
/// and will appear contiguously in the WGSL source.
///
/// SPIR-V (2.11 Structured Control Flow) defines:
///   - loop construct
///   - continue construct
///   - selection construct
/// We also define a "function construct" consisting of all the basic blocks in
/// the function.
///
/// The first block in a construct (by computed block order) is called a
/// "header". For the constructs defined by SPIR-V, the header block is the
/// basic block containing the merge instruction.  The header for the function
/// construct is the entry block of the function.
///
/// Given two constructs A and B, we say "A encloses B" if B is a subset of A,
/// i.e. if every basic block in B is also in A.  Note that a construct encloses
/// itself.
///
/// In a valid SPIR-V module, constructs will nest, meaning given
/// constructs A and B, either A encloses B, or B encloses A, or
/// or they are disjoint (have no basic blocks in commont).
///
/// A loop in a high level language translates into either:
//
///  - a single-block loop, where the loop header branches back to itself.
///     In this case this single-block loop consists only of the *continue
///     construct*.  There is no "loop construct" for this case.
//
///  - a multi-block loop, where the loop back-edge is different from the loop
///     header.
///     This case has both a non-empty loop construct containing at least the
///     loop header, and a non-empty continue construct, containing at least the
///     back-edge block.
///
/// We care about two kinds of selection constructs:
///
///  - if-selection: where the header block ends in OpBranchConditional
///
///  - switch-selection: where the header block ends in OpSwitch
///
struct Construct {
    /// Enumeration for the kinds of structured constructs.
    enum Kind {
        /// The whole function.
        kFunction,
        /// A SPIR-V selection construct, header basic block ending in
        /// OpBrancConditional.
        kIfSelection,
        /// A SPIR-V selection construct, header basic block ending in OpSwitch.
        kSwitchSelection,
        /// A SPIR-V loop construct.
        kLoop,
        /// A SPIR-V continue construct.
        kContinue,
    };

    /// Constructor
    /// @param the_parent parent construct
    /// @param the_depth construct nesting depth
    /// @param the_kind construct kind
    /// @param the_begin_id block id of the first block in the construct
    /// @param the_end_id block id of the first block after the construct, or 0
    /// @param the_begin_pos block order position of the_begin_id
    /// @param the_end_pos block order position of the_end_id or a too-large value
    /// @param the_scope_end_pos block position of the first block past the end of
    /// the WGSL scope
    Construct(const Construct* the_parent,
              int the_depth,
              Kind the_kind,
              uint32_t the_begin_id,
              uint32_t the_end_id,
              uint32_t the_begin_pos,
              uint32_t the_end_pos,
              uint32_t the_scope_end_pos);

    /// @param pos a block position
    /// @returns true if the given block position is inside this construct.
    bool ContainsPos(uint32_t pos) const { return begin_pos <= pos && pos < end_pos; }
    /// Returns true if the given block position is inside the WGSL scope
    /// corresponding to this construct. A loop construct's WGSL scope encloses
    /// the associated continue construct. Otherwise the WGSL scope extent is the
    /// same as the block extent.
    /// @param pos a block position
    /// @returns true if the given block position is inside the WGSL scope.
    bool ScopeContainsPos(uint32_t pos) const { return begin_pos <= pos && pos < scope_end_pos; }

    /// The nearest enclosing construct other than itself, or nullptr if
    /// this construct represents the entire function.
    const Construct* const parent = nullptr;
    /// The nearest enclosing loop construct, if one exists.  Points to `this`
    /// when this is a loop construct.
    const Construct* const enclosing_loop = nullptr;
    /// The nearest enclosing continue construct, if one exists.  Points to
    /// `this` when this is a contnue construct.
    const Construct* const enclosing_continue = nullptr;
    /// The nearest enclosing loop construct or continue construct or
    /// switch-selection construct, if one exists. The signficance is
    /// that a high level language "break" will branch to the merge block
    /// of such an enclosing construct. Points to `this` when this is
    /// a loop construct, a continue construct, or a switch-selection construct.
    const Construct* const enclosing_loop_or_continue_or_switch = nullptr;

    /// Control flow nesting depth. The entry block is at nesting depth 0.
    const int depth = 0;
    /// The construct kind
    const Kind kind = kFunction;
    /// The id of the first block in this structure.
    const uint32_t begin_id = 0;
    /// 0 for kFunction, or the id of the block immediately after this construct
    /// in the computed block order.
    const uint32_t end_id = 0;
    /// The position of block #begin_id in the computed block order.
    const uint32_t begin_pos = 0;
    /// The position of block #end_id in the block order, or the number of
    /// block order elements if #end_id is 0.
    const uint32_t end_pos = 0;
    /// The position of the first block after the WGSL scope corresponding to
    /// this construct.
    const uint32_t scope_end_pos = 0;
};

/// ConstructList is a list of Construct unique pointers.
using ConstructList = std::vector<std::unique_ptr<Construct>>;

/// Converts a construct kind to a string.
/// @param kind the construct kind to convert
/// @returns the string representation
inline std::string ToString(Construct::Kind kind) {
    switch (kind) {
        case Construct::kFunction:
            return "Function";
        case Construct::kIfSelection:
            return "IfSelection";
        case Construct::kSwitchSelection:
            return "SwitchSelection";
        case Construct::kLoop:
            return "Loop";
        case Construct::kContinue:
            return "Continue";
    }
    return "NONE";
}

/// Converts a construct into a short summary string.
/// @param c the construct, which can be null
/// @returns a short summary string
inline std::string ToStringBrief(const Construct* c) {
    if (c) {
        std::stringstream ss;
        ss << ToString(c->kind) << "@" << c->begin_id;
        return ss.str();
    }
    return "null";
}

/// Emits a construct to a stream.
/// @param o the stream
/// @param c the structured construct
/// @returns the stream
inline std::ostream& operator<<(std::ostream& o, const Construct& c) {
    o << "Construct{ " << ToString(c.kind) << " [" << c.begin_pos << "," << c.end_pos << ")"
      << " begin_id:" << c.begin_id << " end_id:" << c.end_id << " depth:" << c.depth;

    o << " parent:" << ToStringBrief(c.parent);

    if (c.scope_end_pos != c.end_pos) {
        o << " scope:[" << c.begin_pos << "," << c.scope_end_pos << ")";
    }

    if (c.enclosing_loop) {
        o << " in-l:" << ToStringBrief(c.enclosing_loop);
    }

    if (c.enclosing_continue) {
        o << " in-c:" << ToStringBrief(c.enclosing_continue);
    }

    if ((c.enclosing_loop_or_continue_or_switch != c.enclosing_loop) &&
        (c.enclosing_loop_or_continue_or_switch != c.enclosing_continue)) {
        o << " in-c-l-s:" << ToStringBrief(c.enclosing_loop_or_continue_or_switch);
    }

    o << " }";
    return o;
}

/// Emits a construct to a stream.
/// @param o the stream
/// @param c the structured construct
/// @returns the stream
inline std::ostream& operator<<(std::ostream& o, const std::unique_ptr<Construct>& c) {
    return o << *(c.get());
}

/// Converts a construct to a string.
/// @param c the construct
/// @returns the string representation
inline std::string ToString(const Construct& c) {
    std::stringstream ss;
    ss << c;
    return ss.str();
}

/// Converts a construct to a string.
/// @param c the construct
/// @returns the string representation
inline std::string ToString(const Construct* c) {
    return c ? ToString(*c) : ToStringBrief(c);
}

/// Converts a unique pointer to a construct to a string.
/// @param c the construct
/// @returns the string representation
inline std::string ToString(const std::unique_ptr<Construct>& c) {
    return ToString(*(c.get()));
}

/// Emits a construct list to a stream.
/// @param o the stream
/// @param cl the construct list
/// @returns the stream
inline std::ostream& operator<<(std::ostream& o, const ConstructList& cl) {
    o << "ConstructList{\n";
    for (const auto& c : cl) {
        o << "  " << c << "\n";
    }
    o << "}";
    return o;
}

/// Converts a construct list to a string.
/// @param cl the construct list
/// @returns the string representation
inline std::string ToString(const ConstructList& cl) {
    std::stringstream ss;
    ss << cl;
    return ss.str();
}

}  // namespace tint::reader::spirv

#endif  // SRC_TINT_READER_SPIRV_CONSTRUCT_H_