validate.cpp

#include "codon/cir/transform/pass.h"
#include "codon/cir/util/irtools.h"
#include "codon/dsl/dsl.h"

using namespace codon::ir;

class ValidateFoo : public transform::OperatorPass {
public:
  static const std::string KEY;
  std::string getKey() const override { return KEY; }

  void handle(AssignInstr *v) override {
    auto *M = v->getModule();
    auto *var = v->getLhs();
    auto *call = cast<CallInstr>(v->getRhs());
    if (!call)
      return;

    auto *foo = util::getFunc(call->getCallee());
    if (!foo || foo->getUnmangledName() != "foo")
      return;

    auto *arg1 = call->front();         // argument of 'foo' call
    auto *arg2 = M->Nr<VarValue>(var);  // result of 'foo' call
    auto *validate =
      M->getOrRealizeFunc("validate", {arg1->getType(), arg2->getType()});
    if (!validate)
      return;

    auto *validateCall = util::call(validate, {arg1, arg2});
    insertAfter(validateCall);  // call 'validate' after 'foo'
  }
};

const std::string ValidateFoo::KEY = "validate-foo";

class MyValidate : public codon::DSL {
public:
  void addIRPasses(transform::PassManager *pm, bool debug) override {
    std::string insertBefore = debug ? "" : "core-folding-pass-group";
    pm->registerPass(std::make_unique<ValidateFoo>(), insertBefore);
  }
};

extern "C" std::unique_ptr<codon::DSL> load() {
  return std::make_unique<MyValidate>();
}