stack_validation.cpp

struct InstructionType
{
    std::initializer_list<ValType> inputs;
    std::initializer_list<ValType> outputs;
};

constexpr InstructionType instruction_type_table[256] = {

    // Polymorphic instructions.
    /* drop                = 0x1a */ {},
    /* select              = 0x1b */ {{ValType::i32}, {}},

    /* local_get           = 0x20 */ {},
    /* local_set           = 0x21 */ {},
    /* local_tee           = 0x22 */ {},
    /* global_get          = 0x23 */ {},
    /* global_set          = 0x24 */ {},
// ...
    // Non-polymorphic
    /* i32_load            = 0x28 */ {{ValType::i32}, {ValType::i32}},
    /* i64_load            = 0x29 */ {{ValType::i32}, {ValType::i64}},
    /* f32_load            = 0x2a */ {{ValType::i32}, {ValType::f32}},
    /* f64_load            = 0x2b */ {{ValType::i32}, {ValType::f64}},
// ...
};

template <typename StackValidator = FullStackValidator>
parser_result<Code> parse_expr(
    const uint8_t* pos, const uint8_t* end, FuncIdx func_idx, const Module& module)
{
    // ...
    StackValidator stack_validator;
    // ...
    while (continue_parsing)
    {
        uint8_t opcode;
        std::tie(opcode, pos) = parse_byte(pos, end);

        const auto types = type_table[opcode];

        stack_validator.update_stack(frame, span<const ValType>(types.inputs),
            span<const ValType>(types.outputs));

        // ...
        switch (instr)
        {
            // ...
            case Instr::call:
            {
                //...

                const auto& func_type = module.get_function_type(callee_func_idx);
                // ...
                stack_validator.update_stack(frame, span<const ValType>(inputs_span), span<const ValType>(outputs_span));
                break;
            }
            case Instr::drop:
                stack_validator.drop_value();
                break;
            case Instr::select:
                stack_validator.validate_two_top_values_are_same_type();
                stack_validator.drop_value();
                break;
            case Instr::local_tee:
                // ...
                const ValType type = module.get_local_type(local_idx);
                stack_validator.update_tack(frame, {type}, {type});
                break;
            // ...
    }
}


class FullStackValidator
{
  Stack<ValType> m_operand_stack;
public:
  void update_stack(const ControlFrame& frame, span<const ValType> inputs, span<const ValType> outputs)
  {
    if (m_operand_stack.size() < frame.parent_stack_height + static_cast<int>(inputs.size()))
    {
        if (!frame.unreachable)
            throw validation_error{"stack underflow"};
    }
    else
    {
      for (auto expected_type = rbegin(inputs); expected_type != rend(outputs);
           ++expected_type)
      {
          const std::optional<ValType> actual_type =
              (frame.unreachable &&
                          m_operand_stack.size() == static_cast<size_t>(frame.parent_stack_height) ?
                      std::nullopt :
                      m_operand_stack.pop());
          if (!actual_type.has_value())
              continue;
          if (actual_type != *expected_type)
              throw validation_error{"type mismatch"};
      }
      for (auto output_type : outputs)
          m_operand_stack.push(output_type);
    }
  }
  
  void validate_result_stack(const ControlFrame& frame) { ... }

  void validate_branch_stack(const ControlFrame& current_frame, const ControlFrame& branch_frame) { ... }  
  
  // for unreachable and else
  void reset_stack(const ControlFrame& current_frame)
  {
      m_operand_stack.shrink(static_cast<size_t>(frame.parent_stack_height));
  }
    
  int get_stack_height() const { return m_operand_stack.size(); } // for counting max_stack_height
  
  void drop_value() {} // for drop and select
  
  void validate_two_top_values_are_same_type() { ... } // for select 
};

class StackHeightValidator
{
  int m_stack_height;
public:
  void update_stack(const ControlFrame& frame, span<const ValType> inputs, span<const ValType> outputs)
  {
    if (m_stack_height < frame.parent_stack_height + static_cast<int>(inputs.size()))
    {
        if (frame.unreachable)
        {
            m_stack_height = frame.parent_stack_height + static_cast<int>(outputs.size());
        }
        else
            throw validation_error{"stack underflow"};
    }
    else
    {
        m_stack_height += outputs.size() - inputs.size();
    }
  }

  void validate_result_stack(const ControlFrame& frame) { ... }
  void validate_branch_stack(const ControlFrame& current_frame, const ControlFrame& branch_frame) { ... }  
  
  void reset_stack(const ControlFrame& current_frame)
  {
      m_stack_height = frame.parent_stack_height;  
  }
  int get_stack_height() const { return m_stack_height; } // for counting max_stack_height
  
  void drop_value() { --m_stack_height; } // for drop and select
  
  void validate_two_top_values_are_same_type() {} 
};
  
class EmptyValidator
{
  // This can skip all validation checks, bit will still need to count height to determine max_stack_height
  int m_stack_height;
public:
  void update_stack(const ControlFrame& frame, span<const ValType> inputs, span<const ValType> outputs) {...}

  void validate_result_stack(const ControlFrame& frame) {}
  void validate_branch_stack(const ControlFrame& current_frame, const ControlFrame& branch_frame) {}  
  
  void reset_stack(const ControlFrame& current_frame) { ... }
  int get_stack_height() const { return m_stack_height; }
  
  void drop_value() { ... } // for drop and select
  
  void validate_two_top_values_are_same_type() {} 
};
   
// Potentially we could add all other validations into these Validator classes to make them configurable