// Helper functions
// Copyright (C) 2024-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/>.

// Dyninst headers
#include <CFG.h>
#include <CodeObject.h>

using Address = Dyninst::Address;
using Block = Dyninst::ParseAPI::Block;
using CodeObject = Dyninst::ParseAPI::CodeObject;
using CodeRegion = Dyninst::ParseAPI::CodeRegion;
using CodeSource = Dyninst::ParseAPI::SymtabCodeSource;

#include <cassert>
#include <cstdlib>
#include <filesystem>
#include <iostream>
#include <set>
#include <string>
#include <utility>

char const*
parse_args (int argc, char const* const* argv)
{
  if (argc == 2)
    return argv[1];
  std::filesystem::path prog {argv[0]};
  std::cerr << "Usage: " << prog.filename ().string () << " EXECUTABLE\n";
  std::exit (1);
}

void
die_for (Address address, std::string const& message)
{
  std::cerr << message << ' ' << std::hex << address << '\n';
  std::exit (1);
}

/// Find next basic block's entry after given address, reparsing if necessary
static Block*
next_block (CodeObject& co, CodeRegion* region, Address address)
{
  auto* blk = co.findBlockByEntry (region, address);
  if (blk != nullptr)
    return blk;
  co.parse (address, true);
  blk = co.findBlockByEntry (region, address);
  return (blk != nullptr) ? blk : co.findNextBlock (region, address);
}

Block*
find_block (CodeSource& cs, CodeObject& co, Address address)
{
  if (!cs.isCode (address))
    die_for (address, "no instruction at");
  std::set <CodeRegion*> regions;
  if (cs.findRegions (address, regions) != 1)
    die_for (address, "not exactly 1 region found for instruction at");
  for (auto* region : regions)
    {
      std::set <Block*> blocks;
      if (co.findBlocks (region, address, blocks) > 0)
        for (auto* blk : blocks) // TODO: choose the best block
          return blk;
      auto* blk = next_block (co, region, region->low ());
      while (blk != nullptr && address > blk->last ())
        blk = next_block (co, region, blk->end ());
      if (blk == nullptr)
        die_for (address, "no block found for instruction at");
      assert (address >= blk->start () && address < blk->end ());
      return blk;
    }
  std::unreachable ();
}