// Function stack size counter
// Copyright (C) 2025  Nguyễn Gia Phong
//
// This file is part of taosc.
//
// Taosc is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// Taosc is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with taosc.  If not, see <https://www.gnu.org/licenses/>.

#include "helpers.hh"

// Dyninst headers
#include <Architecture.h>
#include <CFG.h>
#include <CodeObject.h>
#include <InstructionDecoder.h>
#include <Register.h>

using Architecture = Dyninst::Architecture;
using Address = Dyninst::Address;
using Function = Dyninst::ParseAPI::Function;
using CodeObject = Dyninst::ParseAPI::CodeObject;
using CodeSource = Dyninst::ParseAPI::SymtabCodeSource;
using Instruction = Dyninst::InstructionAPI::Instruction;
using InstructionDecoder = Dyninst::InstructionAPI::InstructionDecoder;
using MachRegister = Dyninst::MachRegister;
using Operand = Dyninst::InstructionAPI::Operand;
using RegisterAST = Dyninst::InstructionAPI::RegisterAST;

#include <cassert>
#include <cstdint>
#include <iostream>
#include <vector>
#include <utility>

static int64_t
immediate (Operand const& operand)
{
  auto const& result = operand.getValue ().get ()->eval ();
  assert (result.defined);
  switch (result.type)
    {
    case Dyninst::InstructionAPI::s8:
    case Dyninst::InstructionAPI::u8:
      return result.val.s8val;
    case Dyninst::InstructionAPI::s16:
    case Dyninst::InstructionAPI::u16:
      return result.val.s16val;
    case Dyninst::InstructionAPI::s32:
    case Dyninst::InstructionAPI::u32:
      return result.val.s32val;
    case Dyninst::InstructionAPI::s64:
    case Dyninst::InstructionAPI::u64:
      return result.val.s64val;
    default:
      std::unreachable ();
    }
}

static int64_t
stack_offset (Instruction const& instruction, Architecture architecture)
{
  switch (instruction.getOperation ().getID ())
    {
    case e_push:
      return Dyninst::getArchAddressWidth (architecture);
    case e_pop:
      return -Dyninst::getArchAddressWidth (architecture);
    case e_sub:
      return immediate (instruction.getOperand (1));
    case e_add:
      return -immediate (instruction.getOperand (1));
    default:
      std::unreachable ();
    }
}

static size_t
stack_offset (char const* buffer, size_t length,
              RegisterAST::Ptr stack_pointer,
              Architecture architecture)
{
  InstructionDecoder decoder {buffer, length, architecture};
  size_t offset = 0;
  for (auto instruction = decoder.decode ();
       instruction.isValid ();
       instruction = decoder.decode ())
    if (instruction.isWritten (stack_pointer))
      offset += stack_offset (instruction, architecture);
  return offset;
}

int
main (int argc, char** argv)
{
  CodeSource cs {parse_args (argc, argv)};
  auto const architecture = cs.getArch ();
  RegisterAST::Ptr const stack_pointer
    {new RegisterAST(MachRegister::getStackPointer (architecture))};
  CodeObject co {&cs};
  co.parse (); // parsed functions have same lifetime as co
  while (!std::cin.eof ())
    {
      Address address;
      std::cin >> std::hex >> address;
      if (std::cin.fail ())
        break;
      auto* const block = find_block (cs, co, address);
      if (block->containingFuncs () < 1)
        die_for (address, "no function found containing instruction at");
      std::vector <Function*> functions;
      block->getFuncs (functions);
      size_t stack_size = 0;
      for (auto* const fun : functions)
        {
          size_t offset = 0;
          for (auto* prologue = fun->entry ();
               prologue;
               prologue = (prologue->targets ().size () == 1
                           ? (*prologue->targets ().begin ())->trg ()
                           : nullptr))
            {
              auto const* const buffer = (char*)
                cs.getPtrToInstruction (prologue->start ());
              offset += stack_offset (buffer, prologue->size (),
                                      stack_pointer, architecture);
            }
          if (offset == 0)
            continue;
          if (stack_size == 0)
            stack_size = offset;
          else if (offset != stack_size)
            die_for (address, "functions with different stack sizes spanning");
        }
      std::cout << stack_size << '\n';
    }
  if (std::cin.eof ())
    return 0;
  std::cerr << "invalid input\n";
  return -1;
}