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/AnalysisUtil.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 = llvm::dyn_cast<component::StructDefOp>(op);
    if (!structDefOp) {
      const char *error_message = "StructAnalysis expects provided op to be a StructDefOp!";
      op->emitError(error_message).report();
      llvm::report_fatal_error(error_message);
    }
    auto maybeModOp = getRootModule(op);
    if (mlir::failed(maybeModOp)) {
      const char *error_message = "StructAnalysis could not find root module from StructDefOp!";
      op->emitError(error_message).report();
      llvm::report_fatal_error(error_message);
    }
    modOp = *maybeModOp;
  }
  virtual ~StructAnalysis() = default;

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

  bool constructed(const Context &ctx) const { return res.contains(ctx); }

  const Result &getResult(const Context &ctx) const {
    ensure(
        constructed(ctx), mlir::Twine(__FUNCTION__) +
                              ": result has not been constructed for struct " +
                              mlir::Twine(getStruct().getName())
    );
    return *res.at(ctx);
  }
 
protected:
  mlir::ModuleOp getModule() const { return modOp; }

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

  void setResult(const Context &ctx, Result &&r) {
    auto [_, inserted] = res.insert(std::make_pair(ctx, std::make_unique<Result>(r)));
    ensure(inserted, "Result already initialized");
  }
 
private:
  mlir::ModuleOp modOp;
  component::StructDefOp structDefOp;
  std::unordered_map<Context, std::unique_ptr<Result>> res;
};
 
template <typename Context>
concept ContextType = requires(const Context &a, const Context &b) {
  { a == b } -> std::convertible_to<bool>;
  { std::hash<Context> {}(a) } -> std::convertible_to<std::size_t>;
};

struct NoContext {};

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

template <typename Result, typename Context, StructAnalysisType<Result, Context> StructAnalysisTy>
class ModuleAnalysis {

  using StructResults = std::map<
      component::StructDefOp, std::reference_wrapper<const Result>,
      NamedOpLocationLess<component::StructDefOp>>;

  using ResultMap = std::unordered_map<Context, StructResults>;
 
public:
  ModuleAnalysis(mlir::Operation *op, const mlir::DataFlowConfig &config = mlir::DataFlowConfig())
      : solver(config) {
    if (modOp = llvm::dyn_cast<mlir::ModuleOp>(op); !modOp) {
      const auto *error_message = "ModuleAnalysis expects provided op to be an mlir::ModuleOp!";
      op->emitError(error_message).report();
      llvm::report_fatal_error(error_message);
    }
  }
  virtual ~ModuleAnalysis() = default;

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

  void ensureAnalysisRun(mlir::AnalysisManager &am) {
    if (!constructed()) {
      runAnalysis(am);
    }
  }

  bool constructed() const { return results.contains(getContext()); }

  bool hasResult(component::StructDefOp op) const {
    return constructed() && results.at(getContext()).contains(op);
  }

  const Result &getResult(component::StructDefOp op) const {
    ensureResultCreated(op);
    return results.at(getContext()).at(op).get();
  }

  const StructResults &getCurrentResults() const {
    ensure(constructed(), "results are not yet constructed for the current context");
    return results.at(getContext());
  }
 
  mlir::DataFlowSolver &getSolver() { return solver; }
 
protected:
  mlir::DataFlowSolver solver;

  virtual void initializeSolver() = 0;

  virtual const Context &getContext() const = 0;

  void constructChildAnalyses(mlir::AnalysisManager &am) {
    auto init = dataflow::loadAndRunRequiredAnalyses(solver, modOp);
    ensure(init.succeeded(), "solver failed to run on module!");
 
    // The analysis is run at the module level so that lattices are computed
    // for global functions as well.
    initializeSolver();
    auto res = solver.initializeAndRun(modOp);
    ensure(res.succeeded(), "solver failed to run on module!");
 
    const Context &ctx = getContext();
    // Force construction of empty results here so `getCurrentResults()` on
    // a module with no inner structs returns no results rather than an assertion
    // failure.
    results[ctx] = {};
    modOp.walk([this, &am, &ctx](component::StructDefOp s) mutable {
      auto &childAnalysis = am.getChildAnalysis<StructAnalysisTy>(s);
      // Don't re-run the analysis if we already have the results.
      // The analysis may have been run as part of a nested analysis.
      if (!childAnalysis.constructed(ctx)) {
        mlir::LogicalResult childAnalysisRes = childAnalysis.runAnalysis(solver, am, ctx);
 
        if (mlir::failed(childAnalysisRes)) {
          auto error_message = "StructAnalysis failed to run for " + s.getName();
          s->emitError(error_message).report();
          llvm::report_fatal_error(error_message);
        }
      }
 
      auto [_, inserted] = results[ctx].insert(
          std::make_pair(s, std::reference_wrapper(childAnalysis.getResult(ctx)))
      );
      ensure(inserted, "struct location conflict");
      return mlir::WalkResult::skip();
    });
  }
 
private:
  mlir::ModuleOp modOp;
  ResultMap results;

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