// Faithful ASAN harness for the luci CircleConst STRING-const OOB read.
//
// This harness reproduces the EXACT read path of
//   compiler/luci/import/src/Nodes/CircleConst.cpp : copy_data<loco::DataType::STRING>
// reached via Importer::importModule -> convert_graph -> CircleConstNodeBuilder::build,
// on the REAL crafted evil.circle parsed by the REAL flatc-generated circle schema.
//
// Faithfulness:
//   * Real schema header (gen_header/circle_schema_generated.h, flatc 24.3.25 over
//     res/CircleSchema/0.10/circle_schema.fbs -- the schema luci_import compiles against).
//   * Real circle::VerifyModelBuffer() gate (same call ImporterEx::importVerifyModule makes).
//   * Real circle::GetModel / subgraph / tensor / buffer navigation.
//   * VectorWrapper<uint8_t> copied verbatim from CircleReader.{h,cpp} (data()/size()).
//   * build()'s buffer-selection + num_elements computation replicated verbatim
//     (CircleConst.cpp:180-251).
//   * copy_data<STRING> read loop copied VERBATIM from CircleConst.cpp:79-96
//     (the over-reading lines), with a minimal CircleConst stub for size<>/at<>.
//
// What is NOT real: the loco/luci IR object graph (CircleConst node) is stubbed, because
// the OOB READ happens entirely on `raw_data` (the flatbuffer Buffer.data) BEFORE any IR
// write. The crashing dereferences are *i32d reads, independent of the IR object.

#include "circle_schema_generated.h"   // real flatc-generated header

#include <cassert>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <fstream>
#include <stdexcept>
#include <string>
#include <vector>

// ---- VectorWrapper<uint8_t>: verbatim from luci/import CircleReader.{h,cpp} ----
namespace luci
{
template <typename T> class VectorWrapper
{
public:
  explicit VectorWrapper(const flatbuffers::Vector<T> *ptr) : _vector(ptr) {}
  const T *data() const { return null() ? nullptr : _vector->data(); }   // CircleReader.cpp:437
  uint32_t size() const { return null() ? 0 : _vector->size(); }         // CircleReader.cpp:432
  bool null() const { return _vector == nullptr; }                       // CircleReader.cpp:469
  bool empty() const { return size() == 0; }                             // CircleReader.cpp:470
private:
  const flatbuffers::Vector<T> *_vector;
};
template <typename T> VectorWrapper<T> wrap(const flatbuffers::Vector<T> *vec)
{
  return VectorWrapper<T>(vec);
}
} // namespace luci

// ---- Minimal CircleConst stub: only the STRING storage methods touched after the read ----
// These run AFTER the OOB read loop; they exist solely so the verbatim copy_data<STRING>
// body compiles. The crash is in the read loop above them.
struct CircleConst
{
  std::vector<std::string> _strings;
  void *sparsityparam() const { return nullptr; }
  template <int DT> void size(uint32_t n) { _strings.resize(n); }
  template <int DT> std::string &at(uint32_t i) { return _strings.at(i); }
};

// ---- copy_data<STRING>: read lines copied VERBATIM from CircleConst.cpp:79-96 ----
// (offset-value validation loop at lines 99-114 runs AFTER and does not prevent the over-read,
//  so for clarity of the crash we keep the exact pre-validation read lines that trigger it.)
static void copy_data_STRING(const luci::VectorWrapper<uint8_t> &raw_data, uint32_t num_elements,
                             CircleConst *const_node)
{
  assert(const_node->sparsityparam() == nullptr);

  const auto *data = reinterpret_cast<const char *>(raw_data.data());      // CircleConst.cpp:81
  const auto *i32d = reinterpret_cast<const int32_t *>(raw_data.data());   // CircleConst.cpp:82

  // de-serialize string data (CircleConst.cpp:84-96)
  assert(static_cast<uint32_t>(*i32d) == num_elements);  // count read  (CircleConst.cpp:88)
  i32d++;                                                 // skip count  (CircleConst.cpp:89)

  std::vector<int32_t> offsets;
  offsets.push_back(*i32d++);                             // offsets[0]  (CircleConst.cpp:92)
  for (uint32_t i = 0; i < num_elements; ++i)
  {
    offsets.push_back(*i32d++);                           // OOB READ    (CircleConst.cpp:95)
  }
  assert(offsets.size() == num_elements + 1);

  (void)data;
  fprintf(stderr, "[harness] survived read loop (no ASAN?) offsets=%zu\n", offsets.size());
}

int main(int argc, char **argv)
{
  const char *path = (argc > 1) ? argv[1] : "evil.circle";

  // Load file into a HEAP allocation, exactly like foder::FileLoader::load() ->
  // std::vector<char> consumed by ImporterEx (ASAN tracks this heap region).
  std::ifstream ifs(path, std::ios::binary | std::ios::ate);
  if (!ifs)
  {
    fprintf(stderr, "cannot open %s\n", path);
    return 2;
  }
  std::streamsize sz = ifs.tellg();
  ifs.seekg(0, std::ios::beg);
  std::vector<char> model_data(static_cast<size_t>(sz));
  if (!ifs.read(model_data.data(), sz))
  {
    fprintf(stderr, "read failed\n");
    return 2;
  }
  auto data_data = reinterpret_cast<const uint8_t *>(model_data.data());
  size_t data_size = model_data.size();
  fprintf(stderr, "[harness] loaded %s (%zu bytes) into heap\n", path, data_size);

  // Real verifier gate (ImporterEx::importVerifyModule).
  {
    flatbuffers::Verifier verifier(data_data, data_size);
    if (!circle::VerifyModelBuffer(verifier))
    {
      fprintf(stderr, "[harness] VerifyModelBuffer FAILED -- file would be rejected\n");
      return 3;
    }
    fprintf(stderr, "[harness] VerifyModelBuffer PASSED (file accepted by importer gate)\n");
  }

  // Importer::importModule -> GetModel
  const circle::Model *model = circle::GetModel(data_data);
  const auto *subgraphs = model->subgraphs();
  if (!subgraphs || subgraphs->size() == 0)
  {
    fprintf(stderr, "no subgraphs\n");
    return 4;
  }
  const circle::SubGraph *sg = subgraphs->Get(0);
  const auto *tensors = sg->tensors();
  const auto *buffers = model->buffers();

  // convert_graph const loop: for each tensor, CircleConstNodeBuilder::build(i)
  for (uint32_t ti = 0; ti < tensors->size(); ++ti)
  {
    const circle::Tensor *const_tensor = tensors->Get(ti);

    // --- build() guards, replicated from CircleConst.cpp:174-257 ---
    if (const_tensor->is_variable())
      continue; // CircleConst.cpp:174
    uint32_t c_buffer = const_tensor->buffer();
    const circle::Buffer *r_buffer = buffers->Get(c_buffer);
    if (r_buffer->offset() == 1 || r_buffer->size() == 1)
      throw std::runtime_error("invalid extended buffer"); // CircleConst.cpp:184

    // We craft an inline buffer (offset==0,size==0), so the else-branch is taken:
    //   buffer = wrap(r_buffer->data());                   (CircleConst.cpp:230)
    luci::VectorWrapper<uint8_t> buffer(nullptr);
    if (r_buffer->offset() > 1)
    {
      fprintf(stderr, "[harness] tensor %u uses extended buffer (not our path)\n", ti);
      continue;
    }
    else
    {
      buffer = luci::wrap(r_buffer->data());
    }

    const auto *shp = const_tensor->shape();
    uint32_t dims = shp ? shp->size() : 0;
    if (dims == 0 && buffer.empty())
      continue; // CircleConst.cpp:233

    // tensoroutputs->find(ti): no operators in evil.circle => never an output. skip check.

    uint32_t num_elements = 1; // CircleConst.cpp:247-251
    for (uint32_t r = 0; r < dims; ++r)
      num_elements = num_elements * shp->Get(r);

    if (buffer.empty() && num_elements > 0)
      continue; // CircleConst.cpp:253

    if (num_elements == 0)
      continue;

    if (const_tensor->type() == circle::TensorType_STRING)
    {
      fprintf(stderr,
              "[harness] tensor %u STRING: num_elements=%u, Buffer.data size=%u bytes; "
              "read loop will touch %u int32 = %u bytes\n",
              ti, num_elements, buffer.size(), num_elements + 1, 4u * (num_elements + 1));
      CircleConst node;
      copy_data_STRING(buffer, num_elements, &node); // <-- CircleConst.cpp:309
    }
  }

  fprintf(stderr, "[harness] done (no crash observed)\n");
  return 0;
}