src/decompiler-ast.inl

/*
 * Copyright 2016 WebAssembly Community Group participants
 *
 * 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.
 */

#include <algorithm>
#include <array>
#include <cassert>
#include <cinttypes>
#include <cstdarg>
#include <cstdio>
#include <iterator>
#include <map>
#include <numeric>
#include <set>
#include <string>
#include <vector>
#include <sstream>

#include "src/cast.h"
#include "src/common.h"
#include "src/expr-visitor.h"
#include "src/ir.h"
#include "src/ir-util.h"
#include "src/literal.h"
#include "src/generate-names.h"
#include "src/stream.h"

namespace wabt {

enum class NodeType {
  FlushToVars,
  FlushedVar,
  Statements,
  EndReturn,
  Decl,
  DeclInit,
  Expr
};

// The AST we're going to convert the standard IR into.
struct Node {
  NodeType ntype;
  ExprType etype;  // Only if ntype == Expr.
  const Expr* e;
  std::vector<Node> children;
  // Node specific annotations.
  union {
    LabelType lt;  // br/br_if target.
    const Var* var;  // Decl/DeclInit.
    struct { Index var_start, var_count; };  // FlushedVar/FlushToVars
  };

  Node()
    : ntype(NodeType::Expr), etype(ExprType::Nop), e(nullptr),
      var(nullptr) {}
  Node(NodeType ntype, ExprType etype, const Expr* e, const Var* v)
    : ntype(ntype), etype(etype), e(e), var(v) {}

  // This value should really never be copied, only moved.
  Node(Node&& rhs) = default;
  Node(const Node& rhs) = delete;
  Node& operator=(Node&& rhs) = default;
  Node& operator=(const Node& rhs) = delete;
};

struct AST {
  AST(ModuleContext& mc, const Func *f) : mc(mc), f(f) {
    if (f) {
      mc.BeginFunc(*f);
      for (Index i = 0; i < f->GetNumParams(); i++) {
        auto name = IndexToAlphaName(i);
        vars_defined.insert(name);
      }
    }
  }

  ~AST() {
    if (f) mc.EndFunc();
  }

  // Create a new node, take nargs existing nodes on the exp stack as children.
  Node& InsertNode(NodeType ntype, ExprType etype, const Expr* e, Index nargs) {
    assert(exp_stack.size() >= nargs);
    Node n { ntype, etype, e, nullptr };
    n.children.reserve(nargs);
    std::move(exp_stack.end() - nargs, exp_stack.end(),
              std::back_inserter(n.children));
    exp_stack.erase(exp_stack.end() - nargs, exp_stack.end());
    exp_stack.push_back(std::move(n));
    return exp_stack.back();
  }

  template<ExprType T> void PreDecl(const VarExpr<T>& ve) {
    exp_stack.emplace(exp_stack.begin() + pre_decl_insertion_point++,
                  NodeType::Decl, ExprType::Nop, nullptr, &ve.var);
  }

  template<ExprType T> void Get(const VarExpr<T>& ve, bool local) {
    if (local && vars_defined.insert(ve.var.name()).second) {
      // Use before def, may happen since locals are guaranteed 0.
      PreDecl(ve);
    }
    InsertNode(NodeType::Expr, T, &ve, 0);
  }

  template<ExprType T> void Set(const VarExpr<T>& ve, bool local) {
    // Seen this var before?
    if (local && vars_defined.insert(ve.var.name()).second) {
      if (value_stack_depth == 1) {
        // Top level, declare it here.
        InsertNode(NodeType::DeclInit, ExprType::Nop, nullptr, 1).var = &ve.var;
        return;
      } else {
        // Inside exp, better leave it as assignment exp and lift the decl out.
        PreDecl(ve);
      }
    }
    InsertNode(NodeType::Expr, T, &ve, 1);
  }

  template<ExprType T> void Block(const BlockExprBase<T>& be, LabelType label) {
    mc.BeginBlock(label, be.block);
    Construct(be.block.exprs, be.block.decl.GetNumResults(), false);
    mc.EndBlock();
    InsertNode(NodeType::Expr, T, &be, 1);
  }

  void Construct(const Expr& e) {
    auto arity = mc.GetExprArity(e);
    switch (e.type()) {
      case ExprType::LocalGet: {
        Get(*cast<LocalGetExpr>(&e), true);
        return;
      }
      case ExprType::GlobalGet: {
        Get(*cast<GlobalGetExpr>(&e), false);
        return;
      }
      case ExprType::LocalSet: {
        Set(*cast<LocalSetExpr>(&e), true);
        return;
      }
      case ExprType::GlobalSet: {
        Set(*cast<GlobalSetExpr>(&e), false);
        return;
      }
      case ExprType::LocalTee: {
        auto& lt = *cast<LocalTeeExpr>(&e);
        Set(lt, true);
        if (value_stack_depth == 1) {  // Tee is the only thing on there.
          Get(lt, true);  // Now Set + Get instead.
        } else {
          // Things are above us on the stack so the Tee can't be eliminated.
          // The Set makes this work as a Tee when consumed by a parent.
        }
        return;
      }
      case ExprType::If: {
        auto ife = cast<IfExpr>(&e);
        value_stack_depth--;  // Condition.
        mc.BeginBlock(LabelType::Block, ife->true_);
        Construct(ife->true_.exprs, ife->true_.decl.GetNumResults(), false);
        if (!ife->false_.empty()) {
          Construct(ife->false_, ife->true_.decl.GetNumResults(), false);
        }
        mc.EndBlock();
        value_stack_depth++;  // Put Condition back.
        InsertNode(NodeType::Expr, ExprType::If, &e, ife->false_.empty() ? 2 : 3);
        return;
      }
      case ExprType::Block: {
        Block(*cast<BlockExpr>(&e), LabelType::Block);
        return;
      }
      case ExprType::Loop: {
        Block(*cast<LoopExpr>(&e), LabelType::Loop);
        return;
      }
      case ExprType::Br: {
        InsertNode(NodeType::Expr, ExprType::Br, &e, 0).lt =
            mc.GetLabel(cast<BrExpr>(&e)->var)->label_type;
        return;
      }
      case ExprType::BrIf: {
        InsertNode(NodeType::Expr, ExprType::BrIf, &e, 1).lt =
            mc.GetLabel(cast<BrIfExpr>(&e)->var)->label_type;
        return;
      }
      default: {
        InsertNode(NodeType::Expr, e.type(), &e, mc.GetExprArity(e).nargs);
        return;
      }
    }
    InsertNode(NodeType::Expr, e.type(), &e, arity.nargs);
  }

  void Construct(const ExprList& es, Index nresults, bool return_results) {
    auto start = exp_stack.size();
    auto value_stack_depth_start = value_stack_depth;
    auto value_stack_in_variables = value_stack_depth;
    bool unreachable = false;
    for (auto& e : es) {
      Construct(e);
      auto arity = mc.GetExprArity(e);
      value_stack_depth -= arity.nargs;
      value_stack_in_variables = std::min(value_stack_in_variables,
                                          value_stack_depth);
      unreachable = unreachable || arity.unreachable;
      assert(unreachable || value_stack_depth >= value_stack_depth_start);
      value_stack_depth += arity.nreturns;
      // We maintain the invariant that a value_stack_depth of N is represented
      // by the last N exp_stack items (each of which returning exactly 1
      // value), all exp_stack items before it return void ("statements").
      // In particular for the wasm stack:
      // - The values between 0 and value_stack_depth_start are part of the
      //   parent block, and not touched here.
      // - The values from there up to value_stack_in_variables are variables
      //   to be used, representing previous statements that flushed the
      //   stack into variables.
      // - Values on top of that up to value_stack_depth are exps returning
      //   a single value.
      // The code below maintains the above invariants. With this in place
      // code "falls into place" the way you expect it.
      if (arity.nreturns != 1) {
        auto num_vars = value_stack_in_variables - value_stack_depth_start;
        auto num_vals = value_stack_depth - value_stack_in_variables;
        auto GenFlushVars = [&](int nargs) {
          auto& ftv = InsertNode(NodeType::FlushToVars, ExprType::Nop, nullptr,
                                 nargs);
          ftv.var_start = flushed_vars;
          ftv.var_count = num_vals;
        };
        auto MoveStatementsBelowVars = [&](size_t amount) {
          std::rotate(exp_stack.end() - num_vars - amount,
                      exp_stack.end() - amount,
                      exp_stack.end());
        };
        auto GenFlushedVars = [&]() {
          // Re-generate these values as vars.
          for (int i = 0; i < num_vals; i++) {
            auto& fv = InsertNode(NodeType::FlushedVar, ExprType::Nop, nullptr,
                                  0);
            fv.var_start = flushed_vars++;
            fv.var_count = 1;
          }
        };
        if (arity.nreturns == 0 &&
            value_stack_depth > value_stack_depth_start) {
          // We have a void item on top of the exp_stack, so we must "lift" the
          // previous values around it.
          // We track what part of the stack is in variables to avoid doing
          // this unnecessarily.
          if (num_vals > 0) {
            // We have actual new values that need lifting.
            // This puts the part of the stack that wasn't already a variable
            // (before the current void exp) into a FlushToVars.
            auto void_exp = std::move(exp_stack.back());
            exp_stack.pop_back();
            GenFlushVars(num_vals);
            exp_stack.push_back(std::move(void_exp));
            // Put this flush statement + void statement before any
            // existing variables.
            MoveStatementsBelowVars(2);
            // Now re-generate these values after the void exp as vars.
            GenFlushedVars();
          } else {
            // We have existing variables that need lifting, but no need to
            // create them anew.
            std::rotate(exp_stack.end() - num_vars - 1, exp_stack.end() - 1,
                        exp_stack.end());
          }
          value_stack_in_variables = value_stack_depth;
        } else if (arity.nreturns > 1) {
          // Multivalue: we flush the stack also.
          // Number of other non-variable values that may be present:
          assert(num_vals >= static_cast<int>(arity.nreturns));
          // Flush multi-value exp + others.
          GenFlushVars(num_vals - arity.nreturns + 1);
          // Put this flush statement before any existing variables.
          MoveStatementsBelowVars(1);
          GenFlushedVars();
          value_stack_in_variables = value_stack_depth;
        }
      } else {
        // Special optimisation: for constant instructions, we can mark these
        // as if they were variables, so they can be re-ordered for free with
        // the above code, without needing new variables!
        // TODO: this would be nice to also do for get_local and maybe others,
        // though that needs a way to ensure there's no set_local in between
        // when they get lifted, so complicates matters a bit.
        if (e.type() == ExprType::Const &&
            value_stack_in_variables == value_stack_depth - 1) {
          value_stack_in_variables++;
        }
      }
    }
    assert(unreachable ||
           value_stack_depth - value_stack_depth_start ==
           static_cast<int>(nresults));
    // Undo any changes to value_stack_depth, since parent takes care of arity
    // changes.
    value_stack_depth = value_stack_depth_start;
    auto end = exp_stack.size();
    assert(end >= start);
    if (return_results && !exp_stack.empty()) {
      if (exp_stack.back().etype == ExprType::Return) {
        if (exp_stack.back().children.empty()) {
          // Return statement at the end of a void function.
          exp_stack.pop_back();
        }
      } else if (nresults) {
        // Combine nresults into a return statement, for when this is used as
        // a function body.
        // TODO: if this is some other kind of block and >1 value is being
        // returned, probably need some kind of syntax to make that clearer.
        InsertNode(NodeType::EndReturn, ExprType::Nop, nullptr, nresults);
      }
    }
    end = exp_stack.size();
    assert(end >= start);
    auto size = end - start;
    if (size != 1) {
      InsertNode(NodeType::Statements, ExprType::Nop, nullptr, size);
    }
  }

  ModuleContext& mc;
  std::vector<Node> exp_stack;
  const Func *f;
  size_t pre_decl_insertion_point = 0;
  int value_stack_depth = 0;
  std::set<std::string> vars_defined;
  Index flushed_vars = 0;
};

}  // namespace wabt