AnalysisWrappers.h

Go to the documentation of this file.

//===-- AnalysisWrappers.h --------------------------------------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2025 Veridise Inc.
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
//===----------------------------------------------------------------------===//
 
#pragma once
 
#include "llzk/Analysis/DenseAnalysis.h"
#include "llzk/Dialect/Struct/IR/Ops.h"
#include "llzk/Util/Compare.h"
#include "llzk/Util/ErrorHelper.h"
#include "llzk/Util/SymbolHelper.h"
 
#include <mlir/IR/BuiltinOps.h>
#include <mlir/Pass/AnalysisManager.h>
 
#include <map>
 
namespace llzk {

template <typename Result, typename Context> class StructAnalysis {
public:
  StructAnalysis(mlir::Operation *op) {
    structDefOp = mlir::dyn_cast<component::StructDefOp>(op);
    if (!structDefOp) {
      auto error_message = "StructAnalysis expects provided op to be a StructDefOp!";
      op->emitError(error_message);
      llvm::report_fatal_error(error_message);
    }
    auto maybeModOp = getRootModule(op);
    if (mlir::failed(maybeModOp)) {
      auto error_message = "StructAnalysis could not find root module from StructDefOp!";
      op->emitError(error_message);
      llvm::report_fatal_error(error_message);
    }
    modOp = *maybeModOp;
  }

  virtual mlir::LogicalResult runAnalysis(
      mlir::DataFlowSolver &solver, mlir::AnalysisManager &moduleAnalysisManager, Context &ctx
  ) = 0;

  bool constructed() const { return res != nullptr; }

  const Result &getResult() const {
    ensure(constructed(), mlir::Twine(__PRETTY_FUNCTION__) + ": result has not been constructed");
    return *res;
  }
 
protected:
  mlir::ModuleOp getModule() const { return modOp; }

  component::StructDefOp getStruct() const { return structDefOp; }

  void setResult(Result &&r) { res = std::make_unique<Result>(r); }
 
private:
  mlir::ModuleOp modOp;
  component::StructDefOp structDefOp;
  std::unique_ptr<Result> res;
};

struct NoContext {};

template <typename Analysis, typename Result, typename Context>
concept StructAnalysisType =
    requires { requires std::is_base_of<StructAnalysis<Result, Context>, Analysis>::value; };

template <typename Result, typename Context, StructAnalysisType<Result, Context> StructAnalysisTy>
class ModuleAnalysis {
  using ResultMap = std::map<
      component::StructDefOp, std::reference_wrapper<const Result>,
      NamedOpLocationLess<component::StructDefOp>>;
 
public:
  ModuleAnalysis(mlir::Operation *op) {
    if (modOp = mlir::dyn_cast<mlir::ModuleOp>(op); !modOp) {
      auto error_message = "ModuleAnalysis expects provided op to be an mlir::ModuleOp!";
      op->emitError(error_message);
      llvm::report_fatal_error(error_message);
    }
  }

  virtual void runAnalysis(mlir::AnalysisManager &am) { constructChildAnalyses(am); }

  bool hasResult(component::StructDefOp op) const { return results.find(op) != results.end(); }

  const Result &getResult(component::StructDefOp op) const {
    ensureResultCreated(op);
    return results.at(op).get();
  }
 
  ResultMap::iterator begin() { return results.begin(); }
  ResultMap::iterator end() { return results.end(); }
  ResultMap::const_iterator cbegin() const { return results.cbegin(); }
  ResultMap::const_iterator cend() const { return results.cend(); }
 
  const mlir::DataFlowSolver &getSolver() const { return solver; }
 
protected:
  virtual void initializeSolver(mlir::DataFlowSolver &solver) = 0;

  virtual Context getContext() = 0;

  void constructChildAnalyses(mlir::AnalysisManager &am) {
    dataflow::markAllOpsAsLive(solver, modOp);
 
    // The analysis is run at the module level so that lattices are computed
    // for global functions as well.
    initializeSolver(solver);
    auto res = solver.initializeAndRun(modOp);
    ensure(res.succeeded(), "solver failed to run on module!");
 
    auto ctx = getContext();
    modOp.walk([this, &am, &ctx](component::StructDefOp s) mutable {
      auto &childAnalysis = am.getChildAnalysis<StructAnalysisTy>(s);
      if (mlir::failed(childAnalysis.runAnalysis(solver, am, ctx))) {
        auto error_message = "StructAnalysis failed to run for " + mlir::Twine(s.getName());
        s->emitError(error_message);
        llvm::report_fatal_error(error_message);
      }
      ensure(results.find(s) == results.end(), "struct location conflict");
      results.insert(std::make_pair(s, std::reference_wrapper(childAnalysis.getResult())));
      return mlir::WalkResult::skip();
    });
  }
 
private:
  mlir::ModuleOp modOp;
  ResultMap results;
  mlir::DataFlowSolver solver;

  void ensureResultCreated(component::StructDefOp op) const {
    ensure(hasResult(op), "Result does not exist for StructDefOp " + mlir::Twine(op.getName()));
  }
};
 
} // namespace llzk