Repository: bpftrace/bpftrace
Pull request #3868: ast: restructure the pass manager
Author: login: amscanne
Created: 2025-03-03T21:34:26Z | Merged: 2025-03-13T19:23:56Z

Stacked PRs:
 * #3888
 * #3885
 * #3869
 * __->__#3868


--- --- ---

### ast: restructure the pass manager


Currently the pass manager does little other than run a sequence of
functions in order. This has led to a few problems:

1) The `BPFtrace` object has logically morphed in a kind of "god object",
   where passes set and mutate global state, which is difficult to
   encapsulate, track or mock.

2) Because of (1), tests for individual passes depend on the implicit
   behavior from all previous passes. Therefore, the same pass plumbing
   is replicated in many places.

Finding myself needing to stash some data resulting from passes
(deserialized bitcode from imports), I decided to take a pass at this
problem rather than adding another field to the AST.

This change introduces a pass context and pass manager which operate
using a `PassContext`. Each pass can require some set of known inputs,
which will be automatically extracted from the `PassContext`.
Furthermore, each pass may produce some known output, which is
automatically injected into to the `PassContext` when the pass is
finished running.

Attempting to add a pass which does not have the required inputs
produced by an earlier pass will generate an error at add time, not at
runtime. Attempting to set an output which has already been set by a
previous pass will also generate an error.

The core idea for this change is that the intermediate state required
(probe attachment info, resources, etc.) will be manifest as these
objects which live in the context. Known passes can produce these
objects, which will then be consumed subsequently.

This makes it easier to design tests around specific passes: rather than
needing to recreate the known universe, you need only manually inject
the specific bits of context that you care about, and validate that the
intended outputs are successfully produced.

Note that this change still leaves the AST and all other intermediate
values as mutable. It would be good to allow for incremental
transformations there as well that could be statically validated, but
this is a much broader task.  Implementing a simple static storage for
passes provides an opportunity to refactor and simplify a lot of the
existing plumbing.

Signed-off-by: Adin Scannell <amscanne@meta.com>



Diff:
diff --git a/src/ast/context.h b/src/ast/context.h
index 9e040472ade4..210adf581f0c 100644
--- a/src/ast/context.h
+++ b/src/ast/context.h
@@ -4,6 +4,7 @@
 #include <vector>
 
 #include "ast/diagnostic.h"
+#include "ast/pass_manager.h"
 
 namespace bpftrace {
 
@@ -42,7 +43,7 @@ class ASTSource {
 // Nodes allocated by an ASTContext will be kept alive for the duration of the
 // owning ASTContext object. The ASTContext also owns the canonical instance of
 // the ASTSource, which is used by the Diagnostics to contextualize errors.
-class ASTContext {
+class ASTContext : public ast::State<"ast"> {
 public:
   ASTContext(std::string &&filename, std::string &&contents);
   ASTContext(const std::string &filename, const std::string &contents);
diff --git a/src/ast/pass_manager.cpp b/src/ast/pass_manager.cpp
index 13ee1e04bd83..fab491cf8254 100644
--- a/src/ast/pass_manager.cpp
+++ b/src/ast/pass_manager.cpp
@@ -1,47 +1,76 @@
 #include "ast/pass_manager.h"
-
-#include <ostream>
-
-#include "ast/passes/printer.h"
-#include "bpftrace.h"
+#include "ast/context.h"
+#include "log.h"
+#include "util/result.h"
 
 namespace bpftrace::ast {
 
-namespace {
-void print(PassContext &ctx, const std::string &name, std::ostream &out)
+std::atomic<int> PassContext::next_type_id_;
+std::unordered_map<int, std::string> PassContext::type_names_;
+
+bool PassContext::ok() const
 {
-  out << "\nAST after: " << name << std::endl;
-  out << "-------------------\n";
-  ast::Printer printer(out);
-  printer.visit(ctx.ast_ctx.root);
-  out << std::endl;
+  int type_id = TypeId<ASTContext>::type_id();
+  if (state_.contains(type_id) || extern_state_.contains(type_id)) {
+    return get<ASTContext>().diagnostics().ok();
+  }
+  return true;
 }
-} // namespace
 
-void PassManager::AddPass(Pass p)
+void PassContext::no_object_failure(int type_id)
 {
-  passes_.push_back(std::move(p));
+  // Rely on the type being available, otherwise how did we get here?
+  LOG(BUG) << "get<" << type_names_[type_id]
+           << "> failed; no object available.";
+  __builtin_unreachable();
 }
 
-int PassManager::Run(PassContext &ctx)
+PassManager &PassManager::add(Pass &&pass)
 {
-  if (bt_debug.find(DebugStage::Ast) != bt_debug.end())
-    print(ctx, "parser", std::cout);
-  for (auto &pass : passes_) {
-    pass.Run(ctx);
-    if (bt_debug.find(DebugStage::Ast) != bt_debug.end())
-      print(ctx, pass.name, std::cout);
-
-    // If the number of errors is non-zero, stop and return.
-    if (!ctx.ast_ctx.diagnostics().ok()) {
-      ctx.ast_ctx.diagnostics().emit(out_);
-      return 1;
+  // Check that the inputs are all available.
+  for (const int type_id : pass.inputs()) {
+    if (!outputs_.contains(type_id)) {
+      auto type_name = PassContext::type_names_[type_id];
+      LOG(BUG)
+          << "Pass " << pass.name() << " requires output " << type_name
+          << ", which is not available; it must be provided by a prior pass.";
     }
   }
+  // Check that the registered output is unique.
+  const int pass_id = passes_.size();
+  for (const int type_id : pass.outputs()) {
+    if (outputs_.contains(type_id)) {
+      auto &orig_pass = passes_[outputs_[type_id]];
+      auto type_name = PassContext::type_names_[type_id];
+      LOG(BUG) << "Pass " << pass.name() << " attempting to register output "
+               << type_name << ", which is already registered by pass "
+               << orig_pass.name() << ".";
+    }
+    // Register the output.
+    outputs_.emplace(type_id, pass_id);
+  }
+  // Add the actual pass.
+  passes_.emplace_back(std::move(pass));
+  return *this;
+}
+
+PassManager &PassManager::add(std::vector<Pass> &&passes)
+{
+  for (auto &pass : passes) {
+    add(std::move(pass));
+  }
+  return *this;
+}
 
-  // Emit all remaining warnings generated by all passes.
-  ctx.ast_ctx.diagnostics().emit(out_);
-  return 0;
+Result<> PassManager::foreach(std::function<Result<>(const Pass &)> fn)
+{
+  for (const auto &pass : passes_) {
+    auto err = fn(pass);
+    if (!err) {
+      return err;
+    }
+  }
+  return OK();
 }
 
 } // namespace bpftrace::ast
diff --git a/src/ast/pass_manager.h b/src/ast/pass_manager.h
index 3f280219ed88..5d7eeef2da0f 100644
--- a/src/ast/pass_manager.h
+++ b/src/ast/pass_manager.h
@@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static_assert(std::is_base_of_v<State, T>,
+                    "type must be derived from State<...>");
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    auto it = state_.find(type_id);
+    if (it != state_.end()) {
+      return *static_cast<T *>(it->second.get());
+    }
+    auto extern_it = extern_state_.find(type_id);
+    if (extern_it != extern_state_.end()) {
+      return static_cast<T &>(extern_it->second.get());
+    }
+    no_object_failure(type_id);
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void no_object_failure(int type_id);
+
+  // put emplaces the value into the map; it must not exist already. In general
+  // this should not be called manually, it should only be called by individual
+  // passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable via
+  // the `type_names_` map. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is an arbitrary pass.
+//
+// It is permitted to produce either no output or a single output, which may be
+// wrapped in `Result`, and may consume zero, one or multiple inputs.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create is the standard Pass constructor.
+  //
+  // For convenience, this template is set up to match against lambdas, and the
+  // partial specialization versions of `createFn` below are used to match
+  // against the resulting std::function. (Of course, a std::function can also
+  // be passed here and the same applies, but using a lambda is encouraged.)
+  template <typename Func>
+  static Pass create(const std::string &name, Func func)
+  {
+    std::function fn(func);
+    return createFn(name, fn);
+  }
 
-  void Run(PassContext &ctx)
+  // run executes the pass.
+  //
+  // The inputs must be present in the context.
+  Result<> run(PassContext &ctx) const
   {
-    fn_(ctx);
-  };
+    return fn_(ctx);
+  }
+
+  const std::string &name() const
+  {
+    return name_;
+  }
+  const std::vector<int> &inputs() const
+  {
+    return inputs_;
+  }
+  const std::vector<int> &outputs() const
+  {
+    return outputs_;
+  }
 
 private:
-  PassFPtr fn_;
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Result<Return>(Inputs &...)> fn)
+  {
+    return createImpl(name, fn);
+  }
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Return(Inputs...)> fn)
+  {
+    // We need to ensure that the arguments are references, but can otherwise
+    // simply wrap this in something that returns OK on the user's behalf.
+    static_assert((std::is_reference_v<Inputs> && ...),
+                  "all arguments must be references");
+    if constexpr (std::is_same_v<Return, void>) {
+      std::function<Result<OK>(Inputs...)> wrap(
+          [fn](Inputs... inputs) -> Result<OK> {
+            fn(inputs...);
+            return OK();
+          });
+      return createImpl(name, wrap);
+    } else {
+      std::function<Result<Return>(Inputs...)> wrap(
+          [fn](Inputs... inputs) -> Result<Return> { return fn(inputs...); });
+      return createImpl(name, wrap);
+    }
+  }
+  template <typename Return, typename... Inputs>
+  static Pass createImpl(const std::string &name,
+                         std::function<Result<Return>(Inputs &...)> fn)
+  {
+    std::vector<int> inputs = { PassContext::TypeId<Inputs>::type_id()... };
+    std::vector<int> outputs;
+    if constexpr (!std::is_same_v<OK, Return>) {
+      // Note that the Return value may be a reference, so we resolve the
+      // `type_id` of the underlying type. The Input types are not references,
+      // because of the way the template is constructed.
+      outputs = { PassContext::TypeId<std::decay_t<Return>>::type_id() };
+    }
 
-public:
-  std::string name;
+    return Pass(
+        name,
+        std::move(inputs),
+        std::move(outputs),
+        [fn](PassContext &ctx) -> Result<> {
+          // Extract a tuple of all input parameters, apply the function
+          // and set the result.
+          auto args = std::make_tuple(std::ref(ctx.get<Inputs>())...);
+          auto result = std::apply(fn, args);
+          if (!result) {
+            return result.takeError();
+          }
+          if constexpr (std::is_same_v<Return, OK>) {
+            return OK();
+          } else {
+            static_assert(
+                std::is_base_of_v<PassContext::State, std::decay_t<Return>>,
+                "return value must be derived from State<...>");
+            if constexpr (std::is_reference_v<Return>) {
+              // Add the external reference.
+              ctx.put(*result);
+            } else {
+              // Just move the result.
+              ctx.put(std::move(*result));
+            }
+            return OK();
+          }
+        });
+  }
+
+  // N.B. external construction is done via `create`.
+  using Fn = std::function<Result<>(PassContext &ctx)>;
+  Pass(std::string name,
+       std::vector<int> &&inputs,
+       std::vector<int> &&outputs,
+       Fn fn)
+      : name_(std::move(name)),
+        inputs_(std::move(inputs)),
+        outputs_(std::move(outputs)),
+        fn_(std::move(fn))
+  {
+  }
+
+  std::string name_;
+  std::vector<int> inputs_;
+  std::vector<int> outputs_;
+  Fn fn_;
 };
 
+// PassManager holds passes in sequence.
+//
+// The PassManager is responsible for verifying the inputs and outputs of the
+// passes ahead of time, so that the produced graph is sound.
 class PassManager {
 public:
-  PassManager(std::ostream &out = std::cerr) : out_(out) {};
+  // registers the given pass, adding to the current list.
+  PassManager &add(Pass &&pass);
 
-  void AddPass(Pass p);
-  [[nodiscard]] int Run(PassContext &ctx);
+  // registers all the given passes.
+  PassManager &add(std::vector<Pass> &&passes);
+
+  // foreach executes the given function on every pass.
+  //
+  // The first pass to encounter errors will result in execution stopping. All
+  // warnings from all passes are collected.
+  Result<> foreach(std::function<Result<>(const Pass &)> fn);
+
+  // put adds a pass which always pushes the given static data.
+  template <typename T>
+  PassManager &put(T &t)
+  {
+    return add(Pass::create(T::type_name(), [&t]() -> T & { return t; }));
+  }
+
+  // run all registered passes.
+  //
+  // This is a canonical example of how to use `foreach` above, in many other
+  // cases a direct call will be suitable.
+  //
+  // Even if the `PassContext` is returned, if an ASTContext is registered, the
+  // user should ensure that no error diagnostics have been produced prior to
+  // extracting state from the context. In this case, it is possible that not
+  // all passes have been run, even if none produced an error.
+  Result<PassContext> run()
+  {
+    PassContext ctx;
+    auto err = foreach([&](auto &pass) -> Result<> {
+      // See above: for convenience, once error diagnostics have been
+      // registered, we skip all remaining passes. This helps to ensure that
+      // users are not overwhelmed by diagnostics, and encodes a common pattern
+      // without the need for explicit error plumbing.
+      if (!ctx.ok())
+        return OK();
+      return pass.run(ctx);
+    });
+    if (!err) {
+      return err.takeError();
+    }
+    return ctx;
+  }
 
 private:
+  // The set of registered passes, in order.
   std::vector<Pass> passes_;
-  std::ostream &out_;
+
+  // The set of registered outputs in aggregate for all other passes. Note that
+  // this is filled in and automatically checked by add, above. The value of
+  // this map is the index into `passes_` above.
+  std::unordered_map<int, int> outputs_;
 };
 
-} // namespace ast
-} // namespace bpftrace
+} // namespace bpftrace::ast
diff --git a/src/ast/passes/config_analyser.cpp b/src/ast/passes/config_analyser.cpp
index 899bfe5d7c7c..2003439230b4 100644
--- a/src/ast/passes/config_analyser.cpp
+++ b/src/ast/passes/config_analyser.cpp
@@ -166,12 +166,12 @@ void ConfigAnalyser::visit(AssignConfigVarStatement &assignment)
 
 Pass CreateConfigPass()
 {
-  auto fn = [](PassContext &ctx) {
-    auto configs = ConfigAnalyser(ctx.b);
-    configs.visit(ctx.ast_ctx.root);
+  auto fn = [](ASTContext &ast, BPFtrace &b) {
+    auto configs = ConfigAnalyser(b);
+    configs.visit(ast.root);
   };
 
-  return { "ConfigAnalyser", fn };
+  return Pass::create("ConfigAnalyser", fn);
 };
 
 } // namespace bpftrace::ast
diff --git a/src/ast/passes/field_analyser.cpp b/src/ast/passes/field_analyser.cpp
index afc1c6c214e7..e5cc1258866e 100644
--- a/src/ast/passes/field_analyser.cpp
+++ b/src/ast/passes/field_analyser.cpp
@@ -322,12 +322,12 @@ void FieldAnalyser::visit(Subprog &subprog)
 
 Pass CreateFieldAnalyserPass()
 {
-  auto fn = [](PassContext &ctx) {
-    FieldAnalyser analyser(ctx.b);
-    analyser.visit(ctx.ast_ctx.root);
+  auto fn = [](ASTContext &ast, BPFtrace &b) {
+    FieldAnalyser analyser(b);
+    analyser.visit(ast.root);
   };
 
-  return { "FieldAnalyser", fn };
+  return Pass::create("FieldAnalyser", fn);
 };
 
 } // namespace bpftrace::ast
diff --git a/src/ast/passes/pid_filter_pass.cpp b/src/ast/passes/pid_filter_pass.cpp
index c58c768a92fc..a05e218cf3e5 100644
--- a/src/ast/passes/pid_filter_pass.cpp
+++ b/src/ast/passes/pid_filter_pass.cpp
@@ -39,19 +39,19 @@ bool probe_needs_pid_filter(AttachPoint *ap)
   return false;
 }
 
-static Statement *create_pid_filter(ASTContext &ctx,
+static Statement *create_pid_filter(ASTContext &ast,
                                     int pid,
                                     const Location &loc)
 {
-  return ctx.make_node<If>(
-      ctx.make_node<Binop>(ctx.make_node<Builtin>("pid", Location(loc)),
+  return ast.make_node<If>(
+      ast.make_node<Binop>(ast.make_node<Builtin>("pid", Location(loc)),
                            Operator::NE,
-                           ctx.make_node<Integer>(pid, Location(loc)),
+                           ast.make_node<Integer>(pid, Location(loc)),
                            Location(loc)),
-      ctx.make_node<Block>(std::vector<Statement *>{ ctx.make_node<Jump>(
+      ast.make_node<Block>(std::vector<Statement *>{ ast.make_node<Jump>(
                                JumpType::RETURN, Location(loc)) },
                            Location(loc)),
-      ctx.make_node<Block>(std::vector<Statement *>{}, Location(loc)),
+      ast.make_node<Block>(std::vector<Statement *>{}, Location(loc)),
       Location(loc));
 }
 
@@ -73,12 +73,10 @@ void PidFilterPass::visit(Probe &probe)
 
 Pass CreatePidFilterPass()
 {
-  auto fn = [](PassContext &ctx) {
-    auto pid_filter = PidFilterPass(ctx.ast_ctx, ctx.b);
-    pid_filter.visit(ctx.ast_ctx.root);
-  };
-
-  return { "PidFilter", fn };
+  return Pass::create("PidFilter", [](ASTContext &ast, BPFtrace &b) {
+    auto pid_filter = PidFilterPass(ast, b);
+    pid_filter.visit(ast.root);
+  });
 };
 
 } // namespace bpftrace::ast
diff --git a/src/ast/passes/portability_analyser.cpp b/src/ast/passes/portability_analyser.cpp
index 9d32d58029e7..8b5c33e68f72 100644
--- a/src/ast/passes/portability_analyser.cpp
+++ b/src/ast/passes/portability_analyser.cpp
@@ -87,17 +87,17 @@ void PortabilityAnalyser::visit(AttachPoint &ap)
 
 Pass CreatePortabilityPass()
 {
-  auto fn = [](PassContext &ctx) {
+  auto fn = [](ASTContext &ast) {
     PortabilityAnalyser analyser;
-    analyser.visit(ctx.ast_ctx.root);
-    if (!ctx.ast_ctx.diagnostics().ok()) {
+    analyser.visit(ast.root);
+    if (!ast.diagnostics().ok()) {
       // Used by runtime test framework to know when to skip an AOT test
       if (std::getenv("__BPFTRACE_NOTIFY_AOT_PORTABILITY_DISABLED"))
         std::cout << "__BPFTRACE_NOTIFY_AOT_PORTABILITY_DISABLED" << std::endl;
     }
   };
 
-  return { "PortabilityAnalyser", fn };
+  return Pass::create("PortabilityAnalyser", fn);
 }
 
 } // namespace bpftrace::ast
diff --git a/src/ast/passes/resource_analyser.cpp b/src/ast/passes/resource_analyser.cpp
index f6fbece02c82..e70be677cd91 100644
--- a/src/ast/passes/resource_analyser.cpp
+++ b/src/ast/passes/resource_analyser.cpp
@@ -491,13 +491,13 @@ void ResourceAnalyser::maybe_allocate_map_key_buffer(const Map &map)
 
 Pass CreateResourcePass()
 {
-  auto fn = [](PassContext &ctx) {
-    ResourceAnalyser analyser(ctx.b);
-    analyser.visit(ctx.ast_ctx.root);
-    ctx.b.resources = analyser.resources();
+  auto fn = [](ASTContext &ast, BPFtrace &b) {
+    ResourceAnalyser analyser(b);
+    analyser.visit(ast.root);
+    b.resources = analyser.resources();
   };
 
-  return { "ResourceAnalyser", fn };
+  return Pass::create("ResourceAnalyser", fn);
 }
 
 } // namespace bpftrace::ast
diff --git a/src/ast/passes/return_path_analyser.cpp b/src/ast/passes/return_path_analyser.cpp
index 44ebc114d32f..8d3e94ed8bf6 100644
--- a/src/ast/passes/return_path_analyser.cpp
+++ b/src/ast/passes/return_path_analyser.cpp
@@ -53,12 +53,12 @@ bool ReturnPathAnalyser::visit(If &if_node)
 
 Pass CreateReturnPathPass()
 {
-  auto fn = [](PassContext &ctx) {
+  auto fn = [](ASTContext &ast) {
     ReturnPathAnalyser return_path;
-    return_path.visit(ctx.ast_ctx.root);
+    return_path.visit(ast.root);
   };
 
-  return { "ReturnPath", fn };
+  return Pass::create("ReturnPath", fn);
 }
 
 } // namespace bpftrace::ast
diff --git a/src/ast/passes/semantic_analyser.cpp b/src/ast/passes/semantic_analyser.cpp
index 280d5372e334..f6043fcfb987 100644
--- a/src/ast/passes/semantic_analyser.cpp
+++ b/src/ast/passes/semantic_analyser.cpp
@@ -4004,12 +4004,12 @@ void SemanticAnalyser::resolve_struct_type(SizedType &type, Node &node)
 
 Pass CreateSemanticPass()
 {
-  auto fn = [](PassContext &ctx) {
-    SemanticAnalyser semantics(ctx.ast_ctx, ctx.b, !ctx.b.cmd_.empty());
+  auto fn = [](ASTContext &ast, BPFtrace &b) {
+    SemanticAnalyser semantics(ast, b, !b.cmd_.empty());
     semantics.analyse();
   };
 
-  return { "Semantic", fn };
+  return Pass::create("Semantic", fn);
 };
 
 Expression *SemanticAnalyser::dereference_if_needed(Expression *expr)
diff --git a/src/bpftrace.h b/src/bpftrace.h
index 26e1af704c7e..0aafde899738 100644
--- a/src/bpftrace.h
+++ b/src/bpftrace.h
@@ -16,6 +16,7 @@
 
 #include "ast/ast.h"
 #include "ast/location.h"
+#include "ast/pass_manager.h"
 #include "attached_probe.h"
 #include "bpfbytecode.h"
 #include "bpffeature.h"
@@ -96,7 +97,7 @@ class WildcardException : public std::exception {
   std::string msg_;
 };
 
-class BPFtrace {
+class BPFtrace : public ast::State<"bpftrace"> {
 public:
   BPFtrace(std::unique_ptr<Output> o = std::make_unique<TextOutput>(std::cout),
            BPFnofeature no_feature = BPFnofeature(),
@@ -111,7 +112,7 @@ class BPFtrace {
         usyms_(*config_)
   {
   }
-  virtual ~BPFtrace();
+  ~BPFtrace() override;
   virtual int add_probe(ast::ASTContext &ctx,
                         const ast::AttachPoint &ap,
                         const ast::Probe &p,
diff --git a/src/main.cpp b/src/main.cpp
index 1286af137d24..2c3402374d8c 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -23,6 +23,7 @@
 #include "ast/passes/field_analyser.h"
 #include "ast/passes/pid_filter_pass.h"
 #include "ast/passes/portability_analyser.h"
+#include "ast/passes/printer.h"
 #include "ast/passes/resource_analyser.h"
 #include "ast/passes/return_path_analyser.h"
 #include "ast/passes/semantic_analyser.h"
@@ -437,23 +438,34 @@ void parse(ast::ASTContext& ast,
   ap_parser.parse();
 }
 
-void CreateDynamicPasses(ast::PassManager& pm)
+void CreateDynamicPasses(std::function<void(ast::Pass&& pass)> add)
 {
-  pm.AddPass(ast::CreateConfigPass());
-  pm.AddPass(ast::CreatePidFilterPass());
-  pm.AddPass(ast::CreateSemanticPass());
-  pm.AddPass(ast::CreateResourcePass());
-  pm.AddPass(ast::CreateReturnPathPass());
+  add(ast::CreateConfigPass());
+  add(ast::CreatePidFilterPass());
+  add(ast::CreateSemanticPass());
+  add(ast::CreateResourcePass());
+  add(ast::CreateReturnPathPass());
 }
 
-void CreateAotPasses(ast::PassManager& pm)
+void CreateAotPasses(std::function<void(ast::Pass&& pass)> add)
 {
-  pm.AddPass(ast::CreateSemanticPass());
-  pm.AddPass(ast::CreatePortabilityPass());
-  pm.AddPass(ast::CreateResourcePass());
-  pm.AddPass(ast::CreateReturnPathPass());
+  add(ast::CreateSemanticPass());
+  add(ast::CreatePortabilityPass());
+  add(ast::CreateResourcePass());
+  add(ast::CreateReturnPathPass());
 }
 
+ast::Pass printPass(const std::string& name)
+{
+  return ast::Pass::create("print-" + name, [=](ast::ASTContext& ast) {
+    std::cerr << "AST after: " << name << std::endl;
+    std::cerr << "-------------------" << std::endl;
+    ast::Printer printer(std::cerr);
+    printer.visit(ast.root);
+    std::cerr << std::endl;
+  });
+};
+
 struct Args {
   std::string pid_str;
   std::string cmd_str;
@@ -974,11 +986,33 @@ int main(int argc, char* argv[])
     return 0;
   }
 
-  ast::PassContext ctx(bpftrace, ast);
+  // Temporarily, we make the full `BPFTrace` object available via the pass
+  // manager (and objects are temporarily mutable). As passes are refactored
+  // into lighter-weight components, the `BPFTrace` object should be decomposed
+  // into its meaningful parts. Furthermore, the codegen and field analysis
+  // passes will be rolled into the pass manager as regular passes; the final
+  // binary is merely one of the outputs that can be extracted.
   ast::PassManager pm;
+  pm.put(ast);
+  pm.put(bpftrace);
+
+  // Wrap all added passes in passes that dump the intermediate state. These
+  // could dump intermediate objects from the context as well, but preserve
+  // existing behavior for now.
+  auto addPass = [&pm](ast::Pass&& pass) {
+    auto name = pass.name();
+    pm.add(std::move(pass));
+    if (bt_debug.find(DebugStage::Ast) != bt_debug.end()) {
+      pm.add(printPass(name));
+    }
+  };
+  if (bt_debug.find(DebugStage::Ast) != bt_debug.end()) {
+    pm.add(printPass("parser"));
+  }
+
   switch (args.build_mode) {
     case BuildMode::DYNAMIC:
-      CreateDynamicPasses(pm);
+      CreateDynamicPasses(addPass);
       break;
     case BuildMode::AHEAD_OF_TIME:
       if (bpftrace.has_dwarf_data()) {
@@ -988,15 +1022,17 @@ int main(int argc, char* argv[])
           std::cout << "__BPFTRACE_NOTIFY_AOT_PORTABILITY_DISABLED"
                     << std::endl;
       }
-      CreateAotPasses(pm);
+      CreateAotPasses(addPass);
       break;
   }
 
   bpftrace.fentry_recursion_check(ast.root);
 
-  auto pmresult = pm.Run(ctx);
-  if (pmresult)
+  auto pmresult = pm.run();
+  if (!pmresult || !ast.diagnostics().ok()) {
+    ast.diagnostics().emit(std::cerr);
     return 1;
+  }
 
   ast::CodegenLLVM llvm(ast, bpftrace);
   BpfBytecode bytecode;
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index e74d244711cd..7e728896dc8b 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -46,6 +46,7 @@ add_executable(bpftrace_test
   procmon.cpp
   probe.cpp
   config_analyser.cpp
+  pass_manager.cpp
   pid_filter_pass.cpp
   resource_analyser.cpp
   result.cpp
diff --git a/tests/pass_manager.cpp b/tests/pass_manager.cpp
new file mode 100644
index 000000000000..17c2e15b122a
--- /dev/null
+++ b/tests/pass_manager.cpp
@@ -0,0 +1,137 @@
+#include "ast/pass_manager.h"
+#include "ast/location.h"
+#include "util/result.h"
+#include "gtest/gtest.h"
+
+using bpftrace::ast::Pass;
+using bpftrace::ast::PassManager;
+
+namespace bpftrace::test::passes {
+
+class Error1 : public ErrorInfo<Error1> {
+public:
+  static char ID;
+  void log(llvm::raw_ostream &) const override
+  {
+  }
+};
+
+class Error2 : public ErrorInfo<Error2> {
+public:
+  static char ID;
+  void log(llvm::raw_ostream &) const override
+  {
+  }
+};
+
+char Error1::ID;
+char Error2::ID;
+
+Pass CreateTest1Pass(bool error = false)
+{
+  return Pass::create("test1", [error]() -> Result<> {
+    if (error) {
+      return make_error<Error1>();
+    }
+    return OK();
+  });
+}
+
+class Test2Output : public ast::State<"test2output"> {};
+
+Pass CreateTest2Pass(bool error = false)
+{
+  return Pass::create("test2", [error]() -> Result<Test2Output> {
+    if (error) {
+      return make_error<Error2>();
+    }
+    return Test2Output();
+  });
+}
+
+class Test3Output : public ast::State<"test3output"> {};
+
+Pass CreateTest3Pass()
+{
+  return Pass::create("test3", [](Test2Output &) -> Result<Test3Output> {
+    return Test3Output();
+  });
+}
+
+TEST(PassManager, noop_pass)
+{
+  PassManager pm;
+  pm.add(Pass::create("void", []() {}));
+  EXPECT_TRUE(bool(pm.run()));
+}
+
+TEST(PassManager, single_pass)
+{
+  PassManager pm;
+  pm.add(CreateTest1Pass());
+  EXPECT_TRUE(bool(pm.run()));
+}
+
+TEST(PassManager, single_pass_with_output)
+{
+  PassManager pm;
+  pm.add(CreateTest2Pass());
+  auto out = pm.run();
+  EXPECT_TRUE(bool(out));
+  out->get<Test2Output>();                   // Should work.
+  EXPECT_DEATH(out->get<Test3Output>(), ""); // Should die.
+}
+
+TEST(PassManager, single_pass_with_error)
+{
+  PassManager pm;
+  pm.add(CreateTest1Pass(true));
+  auto out = pm.run();
+  EXPECT_FALSE(bool(out));
+  EXPECT_TRUE(bool(handleErrors(std::move(out), [](const Error1 &) {})));
+}
+
+TEST(PassManager, multiple_passes)
+{
+  PassManager pm;
+  pm.add(CreateTest1Pass());
+  pm.add(CreateTest1Pass());
+  EXPECT_TRUE(bool(pm.run()));
+}
+
+TEST(PassManager, multiple_passes_with_dependencies)
+{
+  PassManager pm;
+  pm.add(CreateTest2Pass());
+  pm.add(CreateTest3Pass());
+  EXPECT_TRUE(bool(pm.run()));
+}
+
+TEST(PassManager, multiple_passes_with_bad_dependencies)
+{
+  PassManager pm;
+  pm.add(CreateTest1Pass());
+  EXPECT_DEATH(pm.add(CreateTest3Pass()), ""); // Should assert fail.
+}
+
+TEST(PassManager, multiple_passes_with_partial_success)
+{
+  PassManager pm;
+  pm.add(CreateTest1Pass());
+  pm.add(CreateTest2Pass(true));
+  auto out = pm.run();
+  EXPECT_FALSE(bool(out));
+  EXPECT_TRUE(bool(handleErrors(std::move(out), [](const Error2 &) {})));
+}
+
+TEST(PassManager, multiple_passes_complex)
+{
+  PassManager pm;
+  pm.add(CreateTest1Pass());
+  pm.add(CreateTest2Pass());
+  pm.add(CreateTest1Pass());
+  pm.add(CreateTest3Pass());
+  EXPECT_TRUE(bool(pm.run()));
+}
+
+} // namespace bpftrace::test::passes


Reviews:
user: login: danobi
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-12T17:26:15Z
user: login: jordalgo
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-12T21:54:07Z
user: login: amscanne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T08:06:15Z
user: login: amscanne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T08:10:29Z
user: login: amscanne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T08:12:03Z
user: login: amscanne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T08:12:08Z
user: login: amscanne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T08:12:17Z
user: login: amscanne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T08:12:25Z
user: login: amscanne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T08:14:51Z
user: login: amscanne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T08:15:02Z
user: login: amscanne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T08:15:15Z
user: login: amscanne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T08:15:24Z
user: login: amscanne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T08:15:34Z
user: login: jordalgo
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T11:36:47Z
user: login: viktormalik
body: I'm far from fully understanding all the template magic here but it _seems_ not that much complicated so I should be able to get it in a reasonable amount of time if necessary. So, giving my ack.

The overall idea of making the pass infra better is great!
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T16:29:23Z
user: login: danobi
body: nits but lgtm otherwise
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-13T16:49:43Z
user: login: danobi
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-13T16:50:59Z
user: login: amscanne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T18:14:40Z
user: login: amscanne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T18:15:21Z
user: login: amscanne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T18:16:02Z
user: login: amscanne
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-13T18:16:23Z

Review comments:
diff_hunk: @@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    if (state_.contains(type_id)) {
+      auto it = state_.find(type_id);
+      return *static_cast<T *>(it->second.get());
+    } else if (extern_state_.contains(type_id)) {
+      return static_cast<T &>(extern_state_.find(type_id)->second.get());
+    } else {
+      fail(type_id);
+    }
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void fail(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable using
+  // lookup_type below. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Returns the underlying type for this type_id. Note that the type must
+  // exist, or this will have undefined behavior. This is called exclusively
+  // by the pass manager, when statically checking pass types.
+  static const std::string &lookup_type(int type_id)
+  {
+    return type_names_[type_id];
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is a arbitrary pass.
+//
+// It is permitted to produce a single output, but may consume any number of
+// inputs.
+//
+// Note that the inputs and outputs should be statically verified by the pass
+// manager when the pass is registered.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create creates the given pass.
+  //
+  // The passed function may accept an arbitrary set of inputs. Note that the
+  // `type_identity_t` basically makes it so this template will always match,
+  // and the types are resolved later.
+  //
+  // For convenience, this template is set up to match against lambdas, and the
+  // partial specialization that is used below will only work with lambdas.
+  // Furthermore, the expansion will only work with `Result<T>` or `Error`
+  // types, so the lambda will need to return one of those.
+  template <typename Func>
+  static Pass create(const std::string &name, Func func)
+  {
+    std::function fn(func);
+    return createFn(name, fn);
+  }
 
-  void Run(PassContext &ctx)
+  // run executes the pass.
+  //
+  // Before we reach this point, the pass manager will have statically verified
+  // that all outputs are available and that get<Type> will succeed.
+  Result<> run(PassContext &ctx) const
   {
-    fn_(ctx);
-  };
+    return fn_(ctx);
+  }
+
+  const std::string &name() const
+  {
+    return name_;
+  }
+  const std::vector<int> &inputs() const
+  {
+    return inputs_;
+  }
+  const std::vector<int> &outputs() const
+  {
+    return outputs_;
+  }
 
 private:
-  PassFPtr fn_;
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Result<Return>(Inputs &...)> fn)
+  {
+    return createImpl(name, fn);
+  }
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Return(Inputs...)> fn)
path: src/ast/pass_manager.h
user: login: danobi
body: Can't do `Inputs &...` like above? Would be better to do that if possible than to have the static_assert? Or maybe static_assert is better error, as otherwise it'll be some crazy template mismatch error
created_at: 2025-03-12T17:26:15Z
author_association: MEMBER
diff_hunk: @@ -1,47 +1,75 @@
 #include "ast/pass_manager.h"
-
-#include <ostream>
-
-#include "ast/passes/printer.h"
-#include "bpftrace.h"
+#include "ast/context.h"
+#include "log.h"
+#include "util/result.h"
 
 namespace bpftrace::ast {
 
-namespace {
-void print(PassContext &ctx, const std::string &name, std::ostream &out)
+std::atomic<int> PassContext::next_type_id_;
path: src/ast/pass_manager.cpp
user: login: jordalgo
body: JC, do we have multiple threads accessing this or is it just for future proofing?
created_at: 2025-03-12T21:23:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,47 +1,75 @@
 #include "ast/pass_manager.h"
-
-#include <ostream>
-
-#include "ast/passes/printer.h"
-#include "bpftrace.h"
+#include "ast/context.h"
+#include "log.h"
+#include "util/result.h"
 
 namespace bpftrace::ast {
 
-namespace {
-void print(PassContext &ctx, const std::string &name, std::ostream &out)
+std::atomic<int> PassContext::next_type_id_;
+std::unordered_map<int, std::string> PassContext::type_names_;
+
+bool PassContext::ok() const
 {
-  out << "\nAST after: " << name << std::endl;
-  out << "-------------------\n";
-  ast::Printer printer(out);
-  printer.visit(ctx.ast_ctx.root);
-  out << std::endl;
+  int type_id = TypeId<ASTContext>::type_id();
+  if (state_.contains(type_id)) {
+    return get<ASTContext>().diagnostics().ok();
+  }
+  return true;
 }
-} // namespace
 
-void PassManager::AddPass(Pass p)
+void PassContext::fail(int type_id)
 {
-  passes_.push_back(std::move(p));
+  // Rely on the type being available, otherwise how did we get here?
+  LOG(BUG) << "get<" << lookup_type(type_id)
+           << "> failed; no object available.";
+  __builtin_unreachable();
 }
 
-int PassManager::Run(PassContext &ctx)
+PassManager &PassManager::add(Pass &&pass)
 {
-  if (bt_debug.find(DebugStage::Ast) != bt_debug.end())
-    print(ctx, "parser", std::cout);
-  for (auto &pass : passes_) {
-    pass.Run(ctx);
-    if (bt_debug.find(DebugStage::Ast) != bt_debug.end())
-      print(ctx, pass.name, std::cout);
-
-    // If the number of errors is non-zero, stop and return.
-    if (!ctx.ast_ctx.diagnostics().ok()) {
-      ctx.ast_ctx.diagnostics().emit(out_);
-      return 1;
+  // Check that the inputs are all available.
+  for (const int type_id : pass.inputs()) {
+    if (!outputs_.contains(type_id)) {
+      auto type_name = PassContext::lookup_type(type_id);
+      LOG(BUG) << "Pass " << pass.name() << " requires output " << type_name
+               << ", which is not available.";
path: src/ast/pass_manager.cpp
user: login: jordalgo
body: WDYT?
```suggestion
               << ", which is not available. Prior passes need to provide this type as an output";
```
created_at: 2025-03-12T21:25:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,47 +1,75 @@
 #include "ast/pass_manager.h"
-
-#include <ostream>
-
-#include "ast/passes/printer.h"
-#include "bpftrace.h"
+#include "ast/context.h"
+#include "log.h"
+#include "util/result.h"
 
 namespace bpftrace::ast {
 
-namespace {
-void print(PassContext &ctx, const std::string &name, std::ostream &out)
+std::atomic<int> PassContext::next_type_id_;
+std::unordered_map<int, std::string> PassContext::type_names_;
+
+bool PassContext::ok() const
 {
-  out << "\nAST after: " << name << std::endl;
-  out << "-------------------\n";
-  ast::Printer printer(out);
-  printer.visit(ctx.ast_ctx.root);
-  out << std::endl;
+  int type_id = TypeId<ASTContext>::type_id();
+  if (state_.contains(type_id)) {
+    return get<ASTContext>().diagnostics().ok();
+  }
+  return true;
 }
-} // namespace
 
-void PassManager::AddPass(Pass p)
+void PassContext::fail(int type_id)
path: src/ast/pass_manager.cpp
user: login: jordalgo
body: this fail seems a bit specific - what about re-naming to something like `no_object_fail` or just inlining this logic where it's called.
created_at: 2025-03-12T21:40:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    if (state_.contains(type_id)) {
+      auto it = state_.find(type_id);
+      return *static_cast<T *>(it->second.get());
+    } else if (extern_state_.contains(type_id)) {
+      return static_cast<T &>(extern_state_.find(type_id)->second.get());
+    } else {
+      fail(type_id);
+    }
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void fail(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable using
+  // lookup_type below. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Returns the underlying type for this type_id. Note that the type must
+  // exist, or this will have undefined behavior. This is called exclusively
+  // by the pass manager, when statically checking pass types.
+  static const std::string &lookup_type(int type_id)
+  {
+    return type_names_[type_id];
path: src/ast/pass_manager.h
user: login: jordalgo
body: Does this really justify it's own function and comment?
created_at: 2025-03-12T21:44:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    if (state_.contains(type_id)) {
+      auto it = state_.find(type_id);
+      return *static_cast<T *>(it->second.get());
+    } else if (extern_state_.contains(type_id)) {
+      return static_cast<T &>(extern_state_.find(type_id)->second.get());
+    } else {
+      fail(type_id);
+    }
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void fail(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable using
+  // lookup_type below. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Returns the underlying type for this type_id. Note that the type must
+  // exist, or this will have undefined behavior. This is called exclusively
+  // by the pass manager, when statically checking pass types.
+  static const std::string &lookup_type(int type_id)
+  {
+    return type_names_[type_id];
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is a arbitrary pass.
path: src/ast/pass_manager.h
user: login: jordalgo
body: ```suggestion
// Pass is an arbitrary pass.
```
created_at: 2025-03-12T21:45:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    if (state_.contains(type_id)) {
+      auto it = state_.find(type_id);
+      return *static_cast<T *>(it->second.get());
+    } else if (extern_state_.contains(type_id)) {
+      return static_cast<T &>(extern_state_.find(type_id)->second.get());
+    } else {
+      fail(type_id);
+    }
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void fail(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable using
+  // lookup_type below. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Returns the underlying type for this type_id. Note that the type must
+  // exist, or this will have undefined behavior. This is called exclusively
+  // by the pass manager, when statically checking pass types.
+  static const std::string &lookup_type(int type_id)
+  {
+    return type_names_[type_id];
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is a arbitrary pass.
+//
+// It is permitted to produce a single output, but may consume any number of
+// inputs.
path: src/ast/pass_manager.h
user: login: jordalgo
body: ```suggestion
// It is permitted to produce a single output, but may consume multiple inputs.
```
created_at: 2025-03-12T21:46:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    if (state_.contains(type_id)) {
+      auto it = state_.find(type_id);
+      return *static_cast<T *>(it->second.get());
+    } else if (extern_state_.contains(type_id)) {
+      return static_cast<T &>(extern_state_.find(type_id)->second.get());
+    } else {
+      fail(type_id);
+    }
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void fail(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable using
+  // lookup_type below. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Returns the underlying type for this type_id. Note that the type must
+  // exist, or this will have undefined behavior. This is called exclusively
+  // by the pass manager, when statically checking pass types.
+  static const std::string &lookup_type(int type_id)
+  {
+    return type_names_[type_id];
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is a arbitrary pass.
+//
+// It is permitted to produce a single output, but may consume any number of
+// inputs.
+//
+// Note that the inputs and outputs should be statically verified by the pass
+// manager when the pass is registered.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create creates the given pass.
path: src/ast/pass_manager.h
user: login: jordalgo
body: I would have assumed that `create` destroys 😛 
created_at: 2025-03-12T21:47:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    if (state_.contains(type_id)) {
+      auto it = state_.find(type_id);
+      return *static_cast<T *>(it->second.get());
+    } else if (extern_state_.contains(type_id)) {
+      return static_cast<T &>(extern_state_.find(type_id)->second.get());
+    } else {
+      fail(type_id);
+    }
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void fail(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable using
+  // lookup_type below. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Returns the underlying type for this type_id. Note that the type must
+  // exist, or this will have undefined behavior. This is called exclusively
+  // by the pass manager, when statically checking pass types.
+  static const std::string &lookup_type(int type_id)
+  {
+    return type_names_[type_id];
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is a arbitrary pass.
+//
+// It is permitted to produce a single output, but may consume any number of
+// inputs.
+//
+// Note that the inputs and outputs should be statically verified by the pass
+// manager when the pass is registered.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create creates the given pass.
+  //
+  // The passed function may accept an arbitrary set of inputs. Note that the
+  // `type_identity_t` basically makes it so this template will always match,
+  // and the types are resolved later.
+  //
+  // For convenience, this template is set up to match against lambdas, and the
+  // partial specialization that is used below will only work with lambdas.
+  // Furthermore, the expansion will only work with `Result<T>` or `Error`
+  // types, so the lambda will need to return one of those.
+  template <typename Func>
+  static Pass create(const std::string &name, Func func)
+  {
+    std::function fn(func);
+    return createFn(name, fn);
+  }
 
-  void Run(PassContext &ctx)
+  // run executes the pass.
+  //
+  // Before we reach this point, the pass manager will have statically verified
path: src/ast/pass_manager.h
user: login: jordalgo
body: You said this above - any need to repeat it here?
created_at: 2025-03-12T21:48:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    if (state_.contains(type_id)) {
+      auto it = state_.find(type_id);
+      return *static_cast<T *>(it->second.get());
+    } else if (extern_state_.contains(type_id)) {
+      return static_cast<T &>(extern_state_.find(type_id)->second.get());
+    } else {
+      fail(type_id);
+    }
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void fail(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable using
+  // lookup_type below. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Returns the underlying type for this type_id. Note that the type must
+  // exist, or this will have undefined behavior. This is called exclusively
+  // by the pass manager, when statically checking pass types.
+  static const std::string &lookup_type(int type_id)
+  {
+    return type_names_[type_id];
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is a arbitrary pass.
+//
+// It is permitted to produce a single output, but may consume any number of
+// inputs.
+//
+// Note that the inputs and outputs should be statically verified by the pass
+// manager when the pass is registered.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create creates the given pass.
+  //
+  // The passed function may accept an arbitrary set of inputs. Note that the
+  // `type_identity_t` basically makes it so this template will always match,
+  // and the types are resolved later.
+  //
+  // For convenience, this template is set up to match against lambdas, and the
+  // partial specialization that is used below will only work with lambdas.
+  // Furthermore, the expansion will only work with `Result<T>` or `Error`
+  // types, so the lambda will need to return one of those.
+  template <typename Func>
+  static Pass create(const std::string &name, Func func)
+  {
+    std::function fn(func);
+    return createFn(name, fn);
+  }
 
-  void Run(PassContext &ctx)
+  // run executes the pass.
+  //
+  // Before we reach this point, the pass manager will have statically verified
+  // that all outputs are available and that get<Type> will succeed.
+  Result<> run(PassContext &ctx) const
   {
-    fn_(ctx);
-  };
+    return fn_(ctx);
+  }
+
+  const std::string &name() const
+  {
+    return name_;
+  }
+  const std::vector<int> &inputs() const
+  {
+    return inputs_;
+  }
+  const std::vector<int> &outputs() const
+  {
+    return outputs_;
+  }
 
 private:
-  PassFPtr fn_;
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Result<Return>(Inputs &...)> fn)
+  {
+    return createImpl(name, fn);
+  }
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Return(Inputs...)> fn)
+  {
+    // We need to ensure that the arguments are references, but can otherwise
+    // simply wrap this in something that returns OK on the user's behalf.
+    static_assert((std::is_reference_v<Inputs> && ...),
+                  "all arguments must be references");
+    if constexpr (std::is_same_v<Return, void>) {
+      std::function<Result<OK>(Inputs...)> wrap(
+          [fn](Inputs... inputs) -> Result<OK> {
+            fn(inputs...);
+            return OK();
+          });
+      return createImpl(name, wrap);
+    } else {
+      std::function<Result<Return>(Inputs...)> wrap(
+          [fn](Inputs... inputs) -> Result<Return> { return fn(inputs...); });
+      return createImpl(name, wrap);
+    }
+  }
+  template <typename Return, typename... Inputs>
+  static Pass createImpl(const std::string &name,
+                         std::function<Result<Return>(Inputs &...)> fn)
+  {
+    std::vector<int> inputs = { PassContext::TypeId<Inputs>::type_id()... };
+    std::vector<int> outputs;
+    if constexpr (!std::is_same_v<OK, Return>) {
+      // Note that the Return value may be a reference, so we resolve the
+      // `type_id` of the underlying type. The Input types are not references,
+      // because of the way the template is constructed.
+      outputs = { PassContext::TypeId<std::decay_t<Return>>::type_id() };
+    }
 
-public:
-  std::string name;
+    return Pass(name,
+                std::move(inputs),
+                std::move(outputs),
+                [fn](PassContext &ctx) -> Result<> {
+                  // Extract a tuple of all input parameters, apply the function
+                  // and set the result.
+                  auto args = std::make_tuple(std::ref(ctx.get<Inputs>())...);
+                  auto result = std::apply(fn, args);
+                  if (!result) {
+                    return result.takeError();
+                  }
+                  if constexpr (std::is_same_v<Return, OK>) {
+                    return OK();
+                  } else {
+                    if constexpr (std::is_reference_v<Return>) {
+                      // Add the external reference.
+                      ctx.put(*result);
+                    } else {
+                      // Just move the result.
+                      ctx.put(std::move(*result));
+                    }
+                    return OK();
+                  }
+                });
+  }
+
+  // N.B. external construction is done via `create`.
+  using Fn = std::function<Result<>(PassContext &ctx)>;
+  Pass(std::string name,
+       std::vector<int> &&inputs,
+       std::vector<int> &&outputs,
+       Fn fn)
+      : name_(std::move(name)),
+        inputs_(std::move(inputs)),
+        outputs_(std::move(outputs)),
+        fn_(std::move(fn))
+  {
+  }
+
+  std::string name_;
+  std::vector<int> inputs_;
+  std::vector<int> outputs_;
+  Fn fn_;
 };
 
+// PassManager hold passes in sequence.
path: src/ast/pass_manager.h
user: login: jordalgo
body: ```suggestion
// PassManager holds passes in sequence.
```
created_at: 2025-03-12T21:50:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    if (state_.contains(type_id)) {
+      auto it = state_.find(type_id);
+      return *static_cast<T *>(it->second.get());
+    } else if (extern_state_.contains(type_id)) {
+      return static_cast<T &>(extern_state_.find(type_id)->second.get());
+    } else {
+      fail(type_id);
+    }
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void fail(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable using
+  // lookup_type below. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Returns the underlying type for this type_id. Note that the type must
+  // exist, or this will have undefined behavior. This is called exclusively
+  // by the pass manager, when statically checking pass types.
+  static const std::string &lookup_type(int type_id)
+  {
+    return type_names_[type_id];
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is a arbitrary pass.
+//
+// It is permitted to produce a single output, but may consume any number of
+// inputs.
+//
+// Note that the inputs and outputs should be statically verified by the pass
+// manager when the pass is registered.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create creates the given pass.
+  //
+  // The passed function may accept an arbitrary set of inputs. Note that the
+  // `type_identity_t` basically makes it so this template will always match,
+  // and the types are resolved later.
+  //
+  // For convenience, this template is set up to match against lambdas, and the
+  // partial specialization that is used below will only work with lambdas.
+  // Furthermore, the expansion will only work with `Result<T>` or `Error`
+  // types, so the lambda will need to return one of those.
+  template <typename Func>
+  static Pass create(const std::string &name, Func func)
+  {
+    std::function fn(func);
+    return createFn(name, fn);
+  }
 
-  void Run(PassContext &ctx)
+  // run executes the pass.
+  //
+  // Before we reach this point, the pass manager will have statically verified
+  // that all outputs are available and that get<Type> will succeed.
+  Result<> run(PassContext &ctx) const
   {
-    fn_(ctx);
-  };
+    return fn_(ctx);
+  }
+
+  const std::string &name() const
+  {
+    return name_;
+  }
+  const std::vector<int> &inputs() const
+  {
+    return inputs_;
+  }
+  const std::vector<int> &outputs() const
+  {
+    return outputs_;
+  }
 
 private:
-  PassFPtr fn_;
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Result<Return>(Inputs &...)> fn)
+  {
+    return createImpl(name, fn);
+  }
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Return(Inputs...)> fn)
+  {
+    // We need to ensure that the arguments are references, but can otherwise
+    // simply wrap this in something that returns OK on the user's behalf.
+    static_assert((std::is_reference_v<Inputs> && ...),
+                  "all arguments must be references");
+    if constexpr (std::is_same_v<Return, void>) {
+      std::function<Result<OK>(Inputs...)> wrap(
+          [fn](Inputs... inputs) -> Result<OK> {
+            fn(inputs...);
+            return OK();
+          });
+      return createImpl(name, wrap);
+    } else {
+      std::function<Result<Return>(Inputs...)> wrap(
+          [fn](Inputs... inputs) -> Result<Return> { return fn(inputs...); });
+      return createImpl(name, wrap);
+    }
+  }
+  template <typename Return, typename... Inputs>
+  static Pass createImpl(const std::string &name,
+                         std::function<Result<Return>(Inputs &...)> fn)
+  {
+    std::vector<int> inputs = { PassContext::TypeId<Inputs>::type_id()... };
+    std::vector<int> outputs;
+    if constexpr (!std::is_same_v<OK, Return>) {
+      // Note that the Return value may be a reference, so we resolve the
+      // `type_id` of the underlying type. The Input types are not references,
+      // because of the way the template is constructed.
+      outputs = { PassContext::TypeId<std::decay_t<Return>>::type_id() };
+    }
 
-public:
-  std::string name;
+    return Pass(name,
+                std::move(inputs),
+                std::move(outputs),
+                [fn](PassContext &ctx) -> Result<> {
+                  // Extract a tuple of all input parameters, apply the function
+                  // and set the result.
+                  auto args = std::make_tuple(std::ref(ctx.get<Inputs>())...);
+                  auto result = std::apply(fn, args);
+                  if (!result) {
+                    return result.takeError();
+                  }
+                  if constexpr (std::is_same_v<Return, OK>) {
+                    return OK();
+                  } else {
+                    if constexpr (std::is_reference_v<Return>) {
+                      // Add the external reference.
+                      ctx.put(*result);
+                    } else {
+                      // Just move the result.
+                      ctx.put(std::move(*result));
+                    }
+                    return OK();
+                  }
+                });
+  }
+
+  // N.B. external construction is done via `create`.
+  using Fn = std::function<Result<>(PassContext &ctx)>;
+  Pass(std::string name,
+       std::vector<int> &&inputs,
+       std::vector<int> &&outputs,
+       Fn fn)
+      : name_(std::move(name)),
+        inputs_(std::move(inputs)),
+        outputs_(std::move(outputs)),
+        fn_(std::move(fn))
+  {
+  }
+
+  std::string name_;
+  std::vector<int> inputs_;
+  std::vector<int> outputs_;
+  Fn fn_;
 };
 
+// PassManager hold passes in sequence.
+//
+// The PassManager is responsible for verifying the inputs and outputs of the
path: src/ast/pass_manager.h
user: login: jordalgo
body: This is now said a third time - maybe just a summary at the top of PassManager would be better.
created_at: 2025-03-12T21:51:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    if (state_.contains(type_id)) {
+      auto it = state_.find(type_id);
+      return *static_cast<T *>(it->second.get());
+    } else if (extern_state_.contains(type_id)) {
+      return static_cast<T &>(extern_state_.find(type_id)->second.get());
+    } else {
+      fail(type_id);
+    }
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void fail(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable using
+  // lookup_type below. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Returns the underlying type for this type_id. Note that the type must
+  // exist, or this will have undefined behavior. This is called exclusively
+  // by the pass manager, when statically checking pass types.
+  static const std::string &lookup_type(int type_id)
+  {
+    return type_names_[type_id];
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is a arbitrary pass.
+//
+// It is permitted to produce a single output, but may consume any number of
+// inputs.
+//
+// Note that the inputs and outputs should be statically verified by the pass
+// manager when the pass is registered.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create creates the given pass.
+  //
+  // The passed function may accept an arbitrary set of inputs. Note that the
+  // `type_identity_t` basically makes it so this template will always match,
+  // and the types are resolved later.
+  //
+  // For convenience, this template is set up to match against lambdas, and the
+  // partial specialization that is used below will only work with lambdas.
+  // Furthermore, the expansion will only work with `Result<T>` or `Error`
+  // types, so the lambda will need to return one of those.
+  template <typename Func>
+  static Pass create(const std::string &name, Func func)
+  {
+    std::function fn(func);
+    return createFn(name, fn);
+  }
 
-  void Run(PassContext &ctx)
+  // run executes the pass.
+  //
+  // Before we reach this point, the pass manager will have statically verified
+  // that all outputs are available and that get<Type> will succeed.
+  Result<> run(PassContext &ctx) const
   {
-    fn_(ctx);
-  };
+    return fn_(ctx);
+  }
+
+  const std::string &name() const
+  {
+    return name_;
+  }
+  const std::vector<int> &inputs() const
+  {
+    return inputs_;
+  }
+  const std::vector<int> &outputs() const
+  {
+    return outputs_;
+  }
 
 private:
-  PassFPtr fn_;
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Result<Return>(Inputs &...)> fn)
+  {
+    return createImpl(name, fn);
+  }
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Return(Inputs...)> fn)
path: src/ast/pass_manager.h
user: login: amscanne
body: Yeah, I wound up doing this to provide a cleaner error message. I'm sure that I could use concepts and constraints in some complicated way but I found the `static_assert` lets you spell out what's wrong.

Some sample errors:

If you don't pass a reference:
```
In file included from /home/amscanne/bpftrace/tests/pass_manager.cpp:1:
/home/amscanne/bpftrace/src/ast/pass_manager.h: In instantiation of 'static bpftrace::ast::Pass bpftrace::ast::Pass::createFn(const std::string&, std::function<_Res(_ArgTypes ...)>) [with Return = llvm::Expected<bpftrace::test::passes::Test3Output>; Inputs = {bpftrace::test::passes::Test2Output}; std::string = std::__cxx11::basic_string<char>]':
/home/amscanne/bpftrace/src/ast/pass_manager.h:197:20:   required from 'static bpftrace::ast::Pass bpftrace::ast::Pass::create(const std::string&, Func) [with Func = bpftrace::test::passes::CreateTest3Pass()::<lambda(bpftrace::test::passes::Test2Output)>; std::string = std::__cxx11::basic_string<char>]'
  197 |     return createFn(name, fn);
      |            ~~~~~~~~^~~~~~~~~~
/home/amscanne/bpftrace/tests/pass_manager.cpp:56:22:   required from here
   56 |   return Pass::create("test3", [](Test2Output) -> Result<Test3Output> {
      |          ~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
   57 |     return Test3Output();
      |     ~~~~~~~~~~~~~~~~~~~~~
   58 |   });
      |   ~~                  
/home/amscanne/bpftrace/src/ast/pass_manager.h:234:25: error: static assertion failed: all arguments must be references
  234 |     static_assert((std::is_reference_v<Inputs> && ...),
      |                    ~~~~~^~~~~~~~~~~~~~~~~~~~~~
```

If you pass a type that isn't an `ast::State` derived type:
```
In file included from /home/amscanne/bpftrace/tests/pass_manager.cpp:1:
/home/amscanne/bpftrace/src/ast/pass_manager.h: In instantiation of 'static int bpftrace::ast::PassContext::TypeId<T>::type_id() [with T = bpftrace::test::passes::Foo]':
/home/amscanne/bpftrace/src/ast/pass_manager.h:254:69:   required from 'static bpftrace::ast::Pass bpftrace::ast::Pass::createImpl(const std::string&, std::function<llvm::Expected<T>(Inputs& ...)>) [with Return = bpftrace::test::passes::Test3Output; Inputs = {bpftrace::test::passes::Foo}; std::string = std::__cxx11::basic_string<char>]'
  254 |     std::vector<int> inputs = { PassContext::TypeId<Inputs>::type_id()... };
      |                                 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~
/home/amscanne/bpftrace/src/ast/pass_manager.h:227:22:   required from 'static bpftrace::ast::Pass bpftrace::ast::Pass::createFn(const std::string&, std::function<llvm::Expected<T>(Inputs& ...)>) [with Return = bpftrace::test::passes::Test3Output; Inputs = {bpftrace::test::passes::Foo}; std::string = std::__cxx11::basic_string<char>]'
  227 |     return createImpl(name, fn);
      |            ~~~~~~~~~~^~~~~~~~~~
/home/amscanne/bpftrace/src/ast/pass_manager.h:198:20:   required from 'static bpftrace::ast::Pass bpftrace::ast::Pass::create(const std::string&, Func) [with Func = bpftrace::test::passes::CreateTest3Pass()::<lambda(bpftrace::test::passes::Foo&)>; std::string = std::__cxx11::basic_string<char>]'
  198 |     return createFn(name, fn);
      |            ~~~~~~~~^~~~~~~~~~
/home/amscanne/bpftrace/tests/pass_manager.cpp:58:22:   required from here
   58 |   return Pass::create("test3", [](Foo&) -> Result<Test3Output> {
      |          ~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
   59 |     return Test3Output();
      |     ~~~~~~~~~~~~~~~~~~~~~
   60 |   });
      |   ~~                  
/home/amscanne/bpftrace/src/ast/pass_manager.h:53:26: error: static assertion failed: type must be derived from State<...>
   53 |       static_assert(std::is_base_of_v<State, T>);
      |                     ~~~~~^~~~~~~~~~~~~~~~~~~~~~
```

If you try to return an unregistered type:
```
In file included from /home/amscanne/bpftrace/tests/pass_manager.cpp:1:
/home/amscanne/bpftrace/src/ast/pass_manager.h: In instantiation of 'static bpftrace::ast::Pass bpftrace::ast::Pass::createImpl(const std::string&, std::function<llvm::Expected<T>(Inputs& ...)>) [with Return = bpftrace::test::passes::Foo; Inputs = {}; std::string = std::__cxx11::basic_string<char>]':
/home/amscanne/bpftrace/src/ast/pass_manager.h:247:24:   required from 'static bpftrace::ast::Pass bpftrace::ast::Pass::createFn(const std::string&, std::function<_Res(_ArgTypes ...)>) [with Return = bpftrace::test::passes::Foo; Inputs = {}; std::string = std::__cxx11::basic_string<char>]'
  247 |       return createImpl(name, wrap);
      |              ~~~~~~~~~~^~~~~~~~~~~~
/home/amscanne/bpftrace/src/ast/pass_manager.h:198:20:   required from 'static bpftrace::ast::Pass bpftrace::ast::Pass::create(const std::string&, Func) [with Func = bpftrace::test::passes::CreateTest3Pass()::<lambda()>; std::string = std::__cxx11::basic_string<char>]'
  198 |     return createFn(name, fn);
      |            ~~~~~~~~^~~~~~~~~~
/home/amscanne/bpftrace/tests/pass_manager.cpp:58:22:   required from here
   58 |   return Pass::create("test3", []() -> Foo {
      |          ~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~
   59 |     return Foo(); // Test3Output();
      |     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
   60 |   });
      |   ~~                  
/home/amscanne/bpftrace/src/ast/pass_manager.h:277:40: error: static assertion failed: return value must be derived from State<...>
  277 |                     static_assert(std::is_base_of_v<PassContext::State, Return>);
      |                                   ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
```
created_at: 2025-03-13T08:06:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,47 +1,75 @@
 #include "ast/pass_manager.h"
-
-#include <ostream>
-
-#include "ast/passes/printer.h"
-#include "bpftrace.h"
+#include "ast/context.h"
+#include "log.h"
+#include "util/result.h"
 
 namespace bpftrace::ast {
 
-namespace {
-void print(PassContext &ctx, const std::string &name, std::ostream &out)
+std::atomic<int> PassContext::next_type_id_;
path: src/ast/pass_manager.cpp
user: login: amscanne
body: Static initializers have historically been fraught. I know that you have thread-safe guarantees for the *same* static initializer (multiple threads hitting the same static initializer, magically one goes and others block). You *may* have hard guarantees around concurrency for *different* static initializers, but I just prefer to be defensive for this case.
created_at: 2025-03-13T08:10:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,47 +1,75 @@
 #include "ast/pass_manager.h"
-
-#include <ostream>
-
-#include "ast/passes/printer.h"
-#include "bpftrace.h"
+#include "ast/context.h"
+#include "log.h"
+#include "util/result.h"
 
 namespace bpftrace::ast {
 
-namespace {
-void print(PassContext &ctx, const std::string &name, std::ostream &out)
+std::atomic<int> PassContext::next_type_id_;
+std::unordered_map<int, std::string> PassContext::type_names_;
+
+bool PassContext::ok() const
 {
-  out << "\nAST after: " << name << std::endl;
-  out << "-------------------\n";
-  ast::Printer printer(out);
-  printer.visit(ctx.ast_ctx.root);
-  out << std::endl;
+  int type_id = TypeId<ASTContext>::type_id();
+  if (state_.contains(type_id)) {
+    return get<ASTContext>().diagnostics().ok();
+  }
+  return true;
 }
-} // namespace
 
-void PassManager::AddPass(Pass p)
+void PassContext::fail(int type_id)
 {
-  passes_.push_back(std::move(p));
+  // Rely on the type being available, otherwise how did we get here?
+  LOG(BUG) << "get<" << lookup_type(type_id)
+           << "> failed; no object available.";
+  __builtin_unreachable();
 }
 
-int PassManager::Run(PassContext &ctx)
+PassManager &PassManager::add(Pass &&pass)
 {
-  if (bt_debug.find(DebugStage::Ast) != bt_debug.end())
-    print(ctx, "parser", std::cout);
-  for (auto &pass : passes_) {
-    pass.Run(ctx);
-    if (bt_debug.find(DebugStage::Ast) != bt_debug.end())
-      print(ctx, pass.name, std::cout);
-
-    // If the number of errors is non-zero, stop and return.
-    if (!ctx.ast_ctx.diagnostics().ok()) {
-      ctx.ast_ctx.diagnostics().emit(out_);
-      return 1;
+  // Check that the inputs are all available.
+  for (const int type_id : pass.inputs()) {
+    if (!outputs_.contains(type_id)) {
+      auto type_name = PassContext::lookup_type(type_id);
+      LOG(BUG) << "Pass " << pass.name() << " requires output " << type_name
+               << ", which is not available.";
path: src/ast/pass_manager.cpp
user: login: amscanne
body: I've tweaked this message, PTAL.
created_at: 2025-03-13T08:12:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,47 +1,75 @@
 #include "ast/pass_manager.h"
-
-#include <ostream>
-
-#include "ast/passes/printer.h"
-#include "bpftrace.h"
+#include "ast/context.h"
+#include "log.h"
+#include "util/result.h"
 
 namespace bpftrace::ast {
 
-namespace {
-void print(PassContext &ctx, const std::string &name, std::ostream &out)
+std::atomic<int> PassContext::next_type_id_;
+std::unordered_map<int, std::string> PassContext::type_names_;
+
+bool PassContext::ok() const
 {
-  out << "\nAST after: " << name << std::endl;
-  out << "-------------------\n";
-  ast::Printer printer(out);
-  printer.visit(ctx.ast_ctx.root);
-  out << std::endl;
+  int type_id = TypeId<ASTContext>::type_id();
+  if (state_.contains(type_id)) {
+    return get<ASTContext>().diagnostics().ok();
+  }
+  return true;
 }
-} // namespace
 
-void PassManager::AddPass(Pass p)
+void PassContext::fail(int type_id)
path: src/ast/pass_manager.cpp
user: login: amscanne
body: Done.
created_at: 2025-03-13T08:12:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    if (state_.contains(type_id)) {
+      auto it = state_.find(type_id);
+      return *static_cast<T *>(it->second.get());
+    } else if (extern_state_.contains(type_id)) {
+      return static_cast<T &>(extern_state_.find(type_id)->second.get());
+    } else {
+      fail(type_id);
+    }
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void fail(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable using
+  // lookup_type below. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Returns the underlying type for this type_id. Note that the type must
+  // exist, or this will have undefined behavior. This is called exclusively
+  // by the pass manager, when statically checking pass types.
+  static const std::string &lookup_type(int type_id)
+  {
+    return type_names_[type_id];
path: src/ast/pass_manager.h
user: login: amscanne
body: Dropped.
created_at: 2025-03-13T08:12:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    if (state_.contains(type_id)) {
+      auto it = state_.find(type_id);
+      return *static_cast<T *>(it->second.get());
+    } else if (extern_state_.contains(type_id)) {
+      return static_cast<T &>(extern_state_.find(type_id)->second.get());
+    } else {
+      fail(type_id);
+    }
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void fail(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable using
+  // lookup_type below. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Returns the underlying type for this type_id. Note that the type must
+  // exist, or this will have undefined behavior. This is called exclusively
+  // by the pass manager, when statically checking pass types.
+  static const std::string &lookup_type(int type_id)
+  {
+    return type_names_[type_id];
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is a arbitrary pass.
path: src/ast/pass_manager.h
user: login: amscanne
body: Done.
created_at: 2025-03-13T08:12:25Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    if (state_.contains(type_id)) {
+      auto it = state_.find(type_id);
+      return *static_cast<T *>(it->second.get());
+    } else if (extern_state_.contains(type_id)) {
+      return static_cast<T &>(extern_state_.find(type_id)->second.get());
+    } else {
+      fail(type_id);
+    }
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void fail(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable using
+  // lookup_type below. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Returns the underlying type for this type_id. Note that the type must
+  // exist, or this will have undefined behavior. This is called exclusively
+  // by the pass manager, when statically checking pass types.
+  static const std::string &lookup_type(int type_id)
+  {
+    return type_names_[type_id];
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is a arbitrary pass.
+//
+// It is permitted to produce a single output, but may consume any number of
+// inputs.
path: src/ast/pass_manager.h
user: login: amscanne
body: I've tweaked this comment to be more precise regarding inputs and outputs. PTAL.
created_at: 2025-03-13T08:14:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    if (state_.contains(type_id)) {
+      auto it = state_.find(type_id);
+      return *static_cast<T *>(it->second.get());
+    } else if (extern_state_.contains(type_id)) {
+      return static_cast<T &>(extern_state_.find(type_id)->second.get());
+    } else {
+      fail(type_id);
+    }
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void fail(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable using
+  // lookup_type below. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Returns the underlying type for this type_id. Note that the type must
+  // exist, or this will have undefined behavior. This is called exclusively
+  // by the pass manager, when statically checking pass types.
+  static const std::string &lookup_type(int type_id)
+  {
+    return type_names_[type_id];
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is a arbitrary pass.
+//
+// It is permitted to produce a single output, but may consume any number of
+// inputs.
+//
+// Note that the inputs and outputs should be statically verified by the pass
+// manager when the pass is registered.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create creates the given pass.
path: src/ast/pass_manager.h
user: login: amscanne
body: I've clarified that this is the intended constructor.
created_at: 2025-03-13T08:15:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    if (state_.contains(type_id)) {
+      auto it = state_.find(type_id);
+      return *static_cast<T *>(it->second.get());
+    } else if (extern_state_.contains(type_id)) {
+      return static_cast<T &>(extern_state_.find(type_id)->second.get());
+    } else {
+      fail(type_id);
+    }
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void fail(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable using
+  // lookup_type below. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Returns the underlying type for this type_id. Note that the type must
+  // exist, or this will have undefined behavior. This is called exclusively
+  // by the pass manager, when statically checking pass types.
+  static const std::string &lookup_type(int type_id)
+  {
+    return type_names_[type_id];
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is a arbitrary pass.
+//
+// It is permitted to produce a single output, but may consume any number of
+// inputs.
+//
+// Note that the inputs and outputs should be statically verified by the pass
+// manager when the pass is registered.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create creates the given pass.
+  //
+  // The passed function may accept an arbitrary set of inputs. Note that the
+  // `type_identity_t` basically makes it so this template will always match,
+  // and the types are resolved later.
+  //
+  // For convenience, this template is set up to match against lambdas, and the
+  // partial specialization that is used below will only work with lambdas.
+  // Furthermore, the expansion will only work with `Result<T>` or `Error`
+  // types, so the lambda will need to return one of those.
+  template <typename Func>
+  static Pass create(const std::string &name, Func func)
+  {
+    std::function fn(func);
+    return createFn(name, fn);
+  }
 
-  void Run(PassContext &ctx)
+  // run executes the pass.
+  //
+  // Before we reach this point, the pass manager will have statically verified
path: src/ast/pass_manager.h
user: login: amscanne
body: I've dropped two of the three redundant descriptions.
created_at: 2025-03-13T08:15:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    if (state_.contains(type_id)) {
+      auto it = state_.find(type_id);
+      return *static_cast<T *>(it->second.get());
+    } else if (extern_state_.contains(type_id)) {
+      return static_cast<T &>(extern_state_.find(type_id)->second.get());
+    } else {
+      fail(type_id);
+    }
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void fail(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable using
+  // lookup_type below. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Returns the underlying type for this type_id. Note that the type must
+  // exist, or this will have undefined behavior. This is called exclusively
+  // by the pass manager, when statically checking pass types.
+  static const std::string &lookup_type(int type_id)
+  {
+    return type_names_[type_id];
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is a arbitrary pass.
+//
+// It is permitted to produce a single output, but may consume any number of
+// inputs.
+//
+// Note that the inputs and outputs should be statically verified by the pass
+// manager when the pass is registered.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create creates the given pass.
+  //
+  // The passed function may accept an arbitrary set of inputs. Note that the
+  // `type_identity_t` basically makes it so this template will always match,
+  // and the types are resolved later.
+  //
+  // For convenience, this template is set up to match against lambdas, and the
+  // partial specialization that is used below will only work with lambdas.
+  // Furthermore, the expansion will only work with `Result<T>` or `Error`
+  // types, so the lambda will need to return one of those.
+  template <typename Func>
+  static Pass create(const std::string &name, Func func)
+  {
+    std::function fn(func);
+    return createFn(name, fn);
+  }
 
-  void Run(PassContext &ctx)
+  // run executes the pass.
+  //
+  // Before we reach this point, the pass manager will have statically verified
+  // that all outputs are available and that get<Type> will succeed.
+  Result<> run(PassContext &ctx) const
   {
-    fn_(ctx);
-  };
+    return fn_(ctx);
+  }
+
+  const std::string &name() const
+  {
+    return name_;
+  }
+  const std::vector<int> &inputs() const
+  {
+    return inputs_;
+  }
+  const std::vector<int> &outputs() const
+  {
+    return outputs_;
+  }
 
 private:
-  PassFPtr fn_;
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Result<Return>(Inputs &...)> fn)
+  {
+    return createImpl(name, fn);
+  }
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Return(Inputs...)> fn)
+  {
+    // We need to ensure that the arguments are references, but can otherwise
+    // simply wrap this in something that returns OK on the user's behalf.
+    static_assert((std::is_reference_v<Inputs> && ...),
+                  "all arguments must be references");
+    if constexpr (std::is_same_v<Return, void>) {
+      std::function<Result<OK>(Inputs...)> wrap(
+          [fn](Inputs... inputs) -> Result<OK> {
+            fn(inputs...);
+            return OK();
+          });
+      return createImpl(name, wrap);
+    } else {
+      std::function<Result<Return>(Inputs...)> wrap(
+          [fn](Inputs... inputs) -> Result<Return> { return fn(inputs...); });
+      return createImpl(name, wrap);
+    }
+  }
+  template <typename Return, typename... Inputs>
+  static Pass createImpl(const std::string &name,
+                         std::function<Result<Return>(Inputs &...)> fn)
+  {
+    std::vector<int> inputs = { PassContext::TypeId<Inputs>::type_id()... };
+    std::vector<int> outputs;
+    if constexpr (!std::is_same_v<OK, Return>) {
+      // Note that the Return value may be a reference, so we resolve the
+      // `type_id` of the underlying type. The Input types are not references,
+      // because of the way the template is constructed.
+      outputs = { PassContext::TypeId<std::decay_t<Return>>::type_id() };
+    }
 
-public:
-  std::string name;
+    return Pass(name,
+                std::move(inputs),
+                std::move(outputs),
+                [fn](PassContext &ctx) -> Result<> {
+                  // Extract a tuple of all input parameters, apply the function
+                  // and set the result.
+                  auto args = std::make_tuple(std::ref(ctx.get<Inputs>())...);
+                  auto result = std::apply(fn, args);
+                  if (!result) {
+                    return result.takeError();
+                  }
+                  if constexpr (std::is_same_v<Return, OK>) {
+                    return OK();
+                  } else {
+                    if constexpr (std::is_reference_v<Return>) {
+                      // Add the external reference.
+                      ctx.put(*result);
+                    } else {
+                      // Just move the result.
+                      ctx.put(std::move(*result));
+                    }
+                    return OK();
+                  }
+                });
+  }
+
+  // N.B. external construction is done via `create`.
+  using Fn = std::function<Result<>(PassContext &ctx)>;
+  Pass(std::string name,
+       std::vector<int> &&inputs,
+       std::vector<int> &&outputs,
+       Fn fn)
+      : name_(std::move(name)),
+        inputs_(std::move(inputs)),
+        outputs_(std::move(outputs)),
+        fn_(std::move(fn))
+  {
+  }
+
+  std::string name_;
+  std::vector<int> inputs_;
+  std::vector<int> outputs_;
+  Fn fn_;
 };
 
+// PassManager hold passes in sequence.
path: src/ast/pass_manager.h
user: login: amscanne
body: Done.
created_at: 2025-03-13T08:15:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,371 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    if (state_.contains(type_id)) {
+      auto it = state_.find(type_id);
+      return *static_cast<T *>(it->second.get());
+    } else if (extern_state_.contains(type_id)) {
+      return static_cast<T &>(extern_state_.find(type_id)->second.get());
+    } else {
+      fail(type_id);
+    }
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void fail(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable using
+  // lookup_type below. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Returns the underlying type for this type_id. Note that the type must
+  // exist, or this will have undefined behavior. This is called exclusively
+  // by the pass manager, when statically checking pass types.
+  static const std::string &lookup_type(int type_id)
+  {
+    return type_names_[type_id];
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is a arbitrary pass.
+//
+// It is permitted to produce a single output, but may consume any number of
+// inputs.
+//
+// Note that the inputs and outputs should be statically verified by the pass
+// manager when the pass is registered.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create creates the given pass.
+  //
+  // The passed function may accept an arbitrary set of inputs. Note that the
+  // `type_identity_t` basically makes it so this template will always match,
+  // and the types are resolved later.
+  //
+  // For convenience, this template is set up to match against lambdas, and the
+  // partial specialization that is used below will only work with lambdas.
+  // Furthermore, the expansion will only work with `Result<T>` or `Error`
+  // types, so the lambda will need to return one of those.
+  template <typename Func>
+  static Pass create(const std::string &name, Func func)
+  {
+    std::function fn(func);
+    return createFn(name, fn);
+  }
 
-  void Run(PassContext &ctx)
+  // run executes the pass.
+  //
+  // Before we reach this point, the pass manager will have statically verified
+  // that all outputs are available and that get<Type> will succeed.
+  Result<> run(PassContext &ctx) const
   {
-    fn_(ctx);
-  };
+    return fn_(ctx);
+  }
+
+  const std::string &name() const
+  {
+    return name_;
+  }
+  const std::vector<int> &inputs() const
+  {
+    return inputs_;
+  }
+  const std::vector<int> &outputs() const
+  {
+    return outputs_;
+  }
 
 private:
-  PassFPtr fn_;
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Result<Return>(Inputs &...)> fn)
+  {
+    return createImpl(name, fn);
+  }
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Return(Inputs...)> fn)
+  {
+    // We need to ensure that the arguments are references, but can otherwise
+    // simply wrap this in something that returns OK on the user's behalf.
+    static_assert((std::is_reference_v<Inputs> && ...),
+                  "all arguments must be references");
+    if constexpr (std::is_same_v<Return, void>) {
+      std::function<Result<OK>(Inputs...)> wrap(
+          [fn](Inputs... inputs) -> Result<OK> {
+            fn(inputs...);
+            return OK();
+          });
+      return createImpl(name, wrap);
+    } else {
+      std::function<Result<Return>(Inputs...)> wrap(
+          [fn](Inputs... inputs) -> Result<Return> { return fn(inputs...); });
+      return createImpl(name, wrap);
+    }
+  }
+  template <typename Return, typename... Inputs>
+  static Pass createImpl(const std::string &name,
+                         std::function<Result<Return>(Inputs &...)> fn)
+  {
+    std::vector<int> inputs = { PassContext::TypeId<Inputs>::type_id()... };
+    std::vector<int> outputs;
+    if constexpr (!std::is_same_v<OK, Return>) {
+      // Note that the Return value may be a reference, so we resolve the
+      // `type_id` of the underlying type. The Input types are not references,
+      // because of the way the template is constructed.
+      outputs = { PassContext::TypeId<std::decay_t<Return>>::type_id() };
+    }
 
-public:
-  std::string name;
+    return Pass(name,
+                std::move(inputs),
+                std::move(outputs),
+                [fn](PassContext &ctx) -> Result<> {
+                  // Extract a tuple of all input parameters, apply the function
+                  // and set the result.
+                  auto args = std::make_tuple(std::ref(ctx.get<Inputs>())...);
+                  auto result = std::apply(fn, args);
+                  if (!result) {
+                    return result.takeError();
+                  }
+                  if constexpr (std::is_same_v<Return, OK>) {
+                    return OK();
+                  } else {
+                    if constexpr (std::is_reference_v<Return>) {
+                      // Add the external reference.
+                      ctx.put(*result);
+                    } else {
+                      // Just move the result.
+                      ctx.put(std::move(*result));
+                    }
+                    return OK();
+                  }
+                });
+  }
+
+  // N.B. external construction is done via `create`.
+  using Fn = std::function<Result<>(PassContext &ctx)>;
+  Pass(std::string name,
+       std::vector<int> &&inputs,
+       std::vector<int> &&outputs,
+       Fn fn)
+      : name_(std::move(name)),
+        inputs_(std::move(inputs)),
+        outputs_(std::move(outputs)),
+        fn_(std::move(fn))
+  {
+  }
+
+  std::string name_;
+  std::vector<int> inputs_;
+  std::vector<int> outputs_;
+  Fn fn_;
 };
 
+// PassManager hold passes in sequence.
+//
+// The PassManager is responsible for verifying the inputs and outputs of the
path: src/ast/pass_manager.h
user: login: amscanne
body: Dropped two.
created_at: 2025-03-13T08:15:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,366 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static_assert(std::is_base_of_v<State, T>,
+                    "type must be derived from State<...>");
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    auto it = state_.find(type_id);
+    if (it != state_.end()) {
+      return *static_cast<T *>(it->second.get());
+    }
+    auto extern_it = extern_state_.find(type_id);
+    if (extern_it != extern_state_.end()) {
+      return static_cast<T &>(extern_it->second.get());
+    }
+    no_object_failure(type_id);
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void no_object_failure(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable via
+  // the `type_names_` map. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is an arbitrary pass.
+//
+// It is permitted to produce either no output or a single output, which may be
+// wrapped in `Result`, and may consume zero, one or multiple inputs.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create is the standard Pass constructor.
+  //
+  // The passed function may accept an arbitrary set of inputs. Note that the
+  // `type_identity_t` basically makes it so this template will always match,
path: src/ast/pass_manager.h
user: login: danobi
body: type_identity_t? I don't see any further references to it
created_at: 2025-03-13T16:05:29Z
author_association: MEMBER
diff_hunk: @@ -1,64 +1,366 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static_assert(std::is_base_of_v<State, T>,
+                    "type must be derived from State<...>");
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    auto it = state_.find(type_id);
+    if (it != state_.end()) {
+      return *static_cast<T *>(it->second.get());
+    }
+    auto extern_it = extern_state_.find(type_id);
+    if (extern_it != extern_state_.end()) {
+      return static_cast<T &>(extern_it->second.get());
+    }
+    no_object_failure(type_id);
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void no_object_failure(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable via
+  // the `type_names_` map. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is an arbitrary pass.
+//
+// It is permitted to produce either no output or a single output, which may be
+// wrapped in `Result`, and may consume zero, one or multiple inputs.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create is the standard Pass constructor.
+  //
+  // The passed function may accept an arbitrary set of inputs. Note that the
+  // `type_identity_t` basically makes it so this template will always match,
+  // and the types are resolved later.
+  //
+  // For convenience, this template is set up to match against lambdas, and the
+  // partial specialization that is used below will only work with lambdas.
path: src/ast/pass_manager.h
user: login: danobi
body: Not sure I got the "only work with lambdas" bit. Specifically which parts?
created_at: 2025-03-13T16:07:20Z
author_association: MEMBER
diff_hunk: @@ -1,64 +1,366 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static_assert(std::is_base_of_v<State, T>,
+                    "type must be derived from State<...>");
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    auto it = state_.find(type_id);
+    if (it != state_.end()) {
+      return *static_cast<T *>(it->second.get());
+    }
+    auto extern_it = extern_state_.find(type_id);
+    if (extern_it != extern_state_.end()) {
+      return static_cast<T &>(extern_it->second.get());
+    }
+    no_object_failure(type_id);
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void no_object_failure(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
path: src/ast/pass_manager.h
user: login: viktormalik
body: Typo
```suggestion
  // put emplaces the value into the map; it must not exist already. In
```

created_at: 2025-03-13T16:07:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,366 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static_assert(std::is_base_of_v<State, T>,
+                    "type must be derived from State<...>");
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    auto it = state_.find(type_id);
+    if (it != state_.end()) {
+      return *static_cast<T *>(it->second.get());
+    }
+    auto extern_it = extern_state_.find(type_id);
+    if (extern_it != extern_state_.end()) {
+      return static_cast<T &>(extern_it->second.get());
+    }
+    no_object_failure(type_id);
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void no_object_failure(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable via
+  // the `type_names_` map. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is an arbitrary pass.
+//
+// It is permitted to produce either no output or a single output, which may be
+// wrapped in `Result`, and may consume zero, one or multiple inputs.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create is the standard Pass constructor.
+  //
+  // The passed function may accept an arbitrary set of inputs. Note that the
+  // `type_identity_t` basically makes it so this template will always match,
+  // and the types are resolved later.
+  //
+  // For convenience, this template is set up to match against lambdas, and the
+  // partial specialization that is used below will only work with lambdas.
+  // Furthermore, the expansion will only work with `Result<T>` or `Error`
+  // types, so the lambda will need to return one of those.
+  template <typename Func>
+  static Pass create(const std::string &name, Func func)
+  {
+    std::function fn(func);
+    return createFn(name, fn);
+  }
 
-  void Run(PassContext &ctx)
+  // run executes the pass.
+  //
+  // The inputs must be present in the context.
+  Result<> run(PassContext &ctx) const
   {
-    fn_(ctx);
-  };
+    return fn_(ctx);
+  }
+
+  const std::string &name() const
+  {
+    return name_;
+  }
+  const std::vector<int> &inputs() const
+  {
+    return inputs_;
+  }
+  const std::vector<int> &outputs() const
+  {
+    return outputs_;
+  }
 
 private:
-  PassFPtr fn_;
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Result<Return>(Inputs &...)> fn)
+  {
+    return createImpl(name, fn);
+  }
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Return(Inputs...)> fn)
+  {
+    // We need to ensure that the arguments are references, but can otherwise
+    // simply wrap this in something that returns OK on the user's behalf.
+    static_assert((std::is_reference_v<Inputs> && ...),
+                  "all arguments must be references");
+    if constexpr (std::is_same_v<Return, void>) {
+      std::function<Result<OK>(Inputs...)> wrap(
+          [fn](Inputs... inputs) -> Result<OK> {
+            fn(inputs...);
+            return OK();
+          });
+      return createImpl(name, wrap);
+    } else {
+      std::function<Result<Return>(Inputs...)> wrap(
+          [fn](Inputs... inputs) -> Result<Return> { return fn(inputs...); });
+      return createImpl(name, wrap);
+    }
+  }
+  template <typename Return, typename... Inputs>
+  static Pass createImpl(const std::string &name,
+                         std::function<Result<Return>(Inputs &...)> fn)
+  {
+    std::vector<int> inputs = { PassContext::TypeId<Inputs>::type_id()... };
+    std::vector<int> outputs;
+    if constexpr (!std::is_same_v<OK, Return>) {
+      // Note that the Return value may be a reference, so we resolve the
+      // `type_id` of the underlying type. The Input types are not references,
+      // because of the way the template is constructed.
+      outputs = { PassContext::TypeId<std::decay_t<Return>>::type_id() };
+    }
 
-public:
-  std::string name;
+    return Pass(
+        name,
+        std::move(inputs),
+        std::move(outputs),
+        [fn](PassContext &ctx) -> Result<> {
+          // Extract a tuple of all input parameters, apply the function
+          // and set the result.
+          auto args = std::make_tuple(std::ref(ctx.get<Inputs>())...);
+          auto result = std::apply(fn, args);
+          if (!result) {
+            return result.takeError();
+          }
+          if constexpr (std::is_same_v<Return, OK>) {
+            return OK();
+          } else {
+            static_assert(
+                std::is_base_of_v<PassContext::State, std::decay_t<Return>>,
+                "return value must be derived from State<...>");
+            if constexpr (std::is_reference_v<Return>) {
+              // Add the external reference.
+              ctx.put(*result);
+            } else {
+              // Just move the result.
+              ctx.put(std::move(*result));
+            }
+            return OK();
+          }
+        });
+  }
+
+  // N.B. external construction is done via `create`.
+  using Fn = std::function<Result<>(PassContext &ctx)>;
+  Pass(std::string name,
+       std::vector<int> &&inputs,
+       std::vector<int> &&outputs,
+       Fn fn)
+      : name_(std::move(name)),
+        inputs_(std::move(inputs)),
+        outputs_(std::move(outputs)),
+        fn_(std::move(fn))
+  {
+  }
+
+  std::string name_;
+  std::vector<int> inputs_;
+  std::vector<int> outputs_;
+  Fn fn_;
 };
 
+// PassManager holds passes in sequence.
+//
+// The PassManager is responsible for verifying the inputs and outputs of the
+// passes ahead of time, so that the produced graph is sound.
 class PassManager {
 public:
-  PassManager(std::ostream &out = std::cerr) : out_(out) {};
+  // registers the given pass, adding to the current list.
+  PassManager &add(Pass &&pass);
 
-  void AddPass(Pass p);
-  [[nodiscard]] int Run(PassContext &ctx);
+  // registers all the given passes.
+  PassManager &add(std::vector<Pass> &&passes);
+
+  // foreach executes the given function on every pass.
+  //
+  // The first pass to encounter errors will result in execution stopping. All
+  // warnings from all passes are collected.
+  Result<> foreach(std::function<Result<>(const Pass &)> fn);
+
+  // put adds a pass which always pushes the given static data.
+  template <typename T>
+  PassManager &put(T &t)
+  {
+    return add(Pass::create(T::type_name(), [&t]() -> T & { return t; }));
+  }
+
+  // run all registered passes.
+  //
+  // This is a canonical example of how to use `foreach` above, in many other
+  // cases a direct call will be suitable.
+  Result<PassContext> run()
+  {
+    PassContext ctx;
+    auto err = foreach([&](auto &pass) -> Result<> {
+      if (!ctx.ok())
+        return OK(); // Skip the pass.
path: src/ast/pass_manager.h
user: login: danobi
body: ```suggestion
        return OK(); // Previous errors diagnostics, skip the pass.
```
created_at: 2025-03-13T16:28:28Z
author_association: MEMBER
diff_hunk: @@ -1,64 +1,366 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static_assert(std::is_base_of_v<State, T>,
+                    "type must be derived from State<...>");
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    auto it = state_.find(type_id);
+    if (it != state_.end()) {
+      return *static_cast<T *>(it->second.get());
+    }
+    auto extern_it = extern_state_.find(type_id);
+    if (extern_it != extern_state_.end()) {
+      return static_cast<T &>(extern_it->second.get());
+    }
+    no_object_failure(type_id);
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void no_object_failure(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
path: src/ast/pass_manager.h
user: login: amscanne
body: Whoops, fixed.
created_at: 2025-03-13T18:14:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,366 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static_assert(std::is_base_of_v<State, T>,
+                    "type must be derived from State<...>");
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    auto it = state_.find(type_id);
+    if (it != state_.end()) {
+      return *static_cast<T *>(it->second.get());
+    }
+    auto extern_it = extern_state_.find(type_id);
+    if (extern_it != extern_state_.end()) {
+      return static_cast<T &>(extern_it->second.get());
+    }
+    no_object_failure(type_id);
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void no_object_failure(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable via
+  // the `type_names_` map. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is an arbitrary pass.
+//
+// It is permitted to produce either no output or a single output, which may be
+// wrapped in `Result`, and may consume zero, one or multiple inputs.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create is the standard Pass constructor.
+  //
+  // The passed function may accept an arbitrary set of inputs. Note that the
+  // `type_identity_t` basically makes it so this template will always match,
path: src/ast/pass_manager.h
user: login: amscanne
body: Yeah, this got rewritten a few times for sure. The `type_identity_t` was the original name for `TypeId` I think :) 

Anyways, I've simplified this documentation and fixed the below as well. Should now make sense.
created_at: 2025-03-13T18:15:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,366 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static_assert(std::is_base_of_v<State, T>,
+                    "type must be derived from State<...>");
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    auto it = state_.find(type_id);
+    if (it != state_.end()) {
+      return *static_cast<T *>(it->second.get());
+    }
+    auto extern_it = extern_state_.find(type_id);
+    if (extern_it != extern_state_.end()) {
+      return static_cast<T &>(extern_it->second.get());
+    }
+    no_object_failure(type_id);
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void no_object_failure(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable via
+  // the `type_names_` map. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is an arbitrary pass.
+//
+// It is permitted to produce either no output or a single output, which may be
+// wrapped in `Result`, and may consume zero, one or multiple inputs.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create is the standard Pass constructor.
+  //
+  // The passed function may accept an arbitrary set of inputs. Note that the
+  // `type_identity_t` basically makes it so this template will always match,
+  // and the types are resolved later.
+  //
+  // For convenience, this template is set up to match against lambdas, and the
+  // partial specialization that is used below will only work with lambdas.
path: src/ast/pass_manager.h
user: login: amscanne
body: Fixed. It now explains this well I think. It's just "lambdas" encouraged (since otherwise you're doing extra copying)
created_at: 2025-03-13T18:16:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,64 +1,366 @@
 #pragma once
 
+#include <atomic>
 #include <functional>
-#include <iostream>
+#include <memory>
 #include <string>
+#include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
-namespace bpftrace {
+#include "util/result.h"
 
-class BPFtrace;
+namespace bpftrace::ast {
 
-namespace ast {
+// PassContext holds the intermediate state of executing passes.
+//
+// This class is essentially a type-erased generic constant storage.
+class PassContext {
+public:
+  // TypeName is used as a template parameter to work around restrictions on
+  // literals passes as template parameters. This is essentially a string
+  // literal used to name the type.
+  template <size_t N>
+  struct TypeName {
+    constexpr TypeName(const char (&s)[N])
+    {
+      std::copy_n(s, N, value);
+    }
+    char value[N];
+    std::string str() const
+    {
+      return std::string(value, sizeof(value));
+    }
+  };
 
-class ASTContext;
+  // TypeId is a class which, when specialized, returns the type_id. This
+  // works by storing a static variable for each specialization of the class,
+  // which effectively memoizes an ID binding for the specific class.
+  template <typename T>
+  struct TypeId {
+    static int type_id()
+    {
+      // Register by dispatching via `type_name` for the passed type. This is
+      // all going to be devirtualized most likely, but this does enforce that
+      // the class passed at least implements the static method, which is
+      // really all we can ask for.
+      //
+      // Note that the `static` here is doing all the heavy lifting; this is
+      // ensuring that the compiler is linking this symbol at most once, and it
+      // will be initialized at most once, giving us a unique identifier for
+      // this particular type.
+      static_assert(std::is_base_of_v<State, T>,
+                    "type must be derived from State<...>");
+      static int type_id = PassContext::add_type(T::type_name());
+      return type_id;
+    }
+  };
 
-// Context/config for passes
-//
-// Note: Most state should end up in the BPFtrace class instead of here
+  // State is an object which is the result of some specific pass.
+  //
+  // Passes are allowed to mutate the AST, and produce intermediate results
+  // that may be consumed by other passes. Each of these results must be a
+  // State, and will be immutable once it is produced.
+  class State {
+  public:
+    virtual ~State() = default;
+  };
 
-struct PassContext {
-public:
-  PassContext(BPFtrace &b, ASTContext &ast_ctx) : b(b), ast_ctx(ast_ctx) {};
-  BPFtrace &b;
-  ASTContext &ast_ctx;
+  // get fetches the given type; which must have been previously stored.
+  //
+  // N.B. In the future, these objects should be made immutable. Right now,
+  // this is preserved so that we can pass through the `BPFTrace` object which
+  // is consumed by *most* passes in a read-write fashion.
+  template <typename T>
+  T &get() const
+  {
+    int type_id = TypeId<T>::type_id();
+    auto it = state_.find(type_id);
+    if (it != state_.end()) {
+      return *static_cast<T *>(it->second.get());
+    }
+    auto extern_it = extern_state_.find(type_id);
+    if (extern_it != extern_state_.end()) {
+      return static_cast<T &>(extern_it->second.get());
+    }
+    no_object_failure(type_id);
+  }
+
+private:
+  // for the failed lookup path, see above.
+  [[noreturn]] static void no_object_failure(int type_id);
+
+  // put emplaces the the value into the map; it must not exist already. In
+  // general this should not be called manually, it should only be called by
+  // individual passes during execution (the `run` function below).
+  template <typename T>
+  void put(T &&t)
+  {
+    int type_id = TypeId<T>::type_id();
+    state_.emplace(type_id, std::make_unique<T>(std::move(t)));
+  }
+
+  // put can also provide an external reference for the map; its lifecycle is
+  // managed outside the scope of the passes. While passes can produce results
+  // that are references, this should be generally discouraged (and this
+  // capability should eventually be removed). Like `put`, this should not be
+  // called directly and should only be called by `run`, below.
+  template <typename T>
+  void put(T &t)
+  {
+    int type_id = TypeId<T>::type_id();
+    extern_state_.emplace(type_id, t);
+  }
+
+  // Will register the given type id and the name, which is resolvable via
+  // the `type_names_` map. This is called exclusively by `type_id` above.
+  static int add_type(std::string &&type_name)
+  {
+    int type_id = next_type_id_++;
+    type_names_[type_id] = std::move(type_name);
+    return type_id;
+  }
+
+  // Provide a short-hand to determine:
+  // - Whether the context contains an `ASTContext`, and,
+  // - Whether this `ASTContext` has error diagnostics.
+  //
+  // If both of the above are true, then false is returned. This is a
+  // convenience method provided for the `PassManager` below, as passes will
+  // stop when an error diagnostic has been generated.
+  bool ok() const;
+
+  static std::atomic<int> next_type_id_;
+  static std::unordered_map<int, std::string> type_names_;
+
+  std::unordered_map<int, std::unique_ptr<State>> state_;
+  std::unordered_map<int, std::reference_wrapper<State>> extern_state_;
+
+  friend class PassManager;
+  friend class Pass;
 };
 
-using PassFPtr = std::function<void(PassContext &)>;
+// State is the template class used for states that may be stored in a
+// PassContext. In order to consume or produce information in a given pass,
+// simply override this class with the desired type.
+//
+// For example:
+//
+//   class Variables : State<"variables"> {
+//     // my fields
+//   };
+//
+// The `Variables` class can then be stored in the pass context.
+template <PassContext::TypeName name>
+class State : public PassContext::State {
+public:
+  static std::string type_name()
+  {
+    return name.str();
+  }
+};
 
-// Base pass
+// Pass is an arbitrary pass.
+//
+// It is permitted to produce either no output or a single output, which may be
+// wrapped in `Result`, and may consume zero, one or multiple inputs.
+//
+// In general, this type should be constructed automatically, e.g.:
+//
+//    Pass CreateFooPass() {
+//      return Pass::create("mypass", [](MyInput &input) {
+//        return MyOutput();
+//      });
+//    }
+//
+// Pass functions must already return an `Result` instance.
 class Pass {
 public:
-  Pass() = delete;
-  Pass(std::string name, PassFPtr fn)
-      : fn_(std::move(fn)), name(std::move(name)) {};
+  Pass(const Pass &) = delete;
+  Pass(Pass &&) = default;
+  Pass &operator=(const Pass &) = delete;
+  Pass &operator=(Pass &&) = default;
 
-  virtual ~Pass() = default;
+  // create is the standard Pass constructor.
+  //
+  // The passed function may accept an arbitrary set of inputs. Note that the
+  // `type_identity_t` basically makes it so this template will always match,
+  // and the types are resolved later.
+  //
+  // For convenience, this template is set up to match against lambdas, and the
+  // partial specialization that is used below will only work with lambdas.
+  // Furthermore, the expansion will only work with `Result<T>` or `Error`
+  // types, so the lambda will need to return one of those.
+  template <typename Func>
+  static Pass create(const std::string &name, Func func)
+  {
+    std::function fn(func);
+    return createFn(name, fn);
+  }
 
-  void Run(PassContext &ctx)
+  // run executes the pass.
+  //
+  // The inputs must be present in the context.
+  Result<> run(PassContext &ctx) const
   {
-    fn_(ctx);
-  };
+    return fn_(ctx);
+  }
+
+  const std::string &name() const
+  {
+    return name_;
+  }
+  const std::vector<int> &inputs() const
+  {
+    return inputs_;
+  }
+  const std::vector<int> &outputs() const
+  {
+    return outputs_;
+  }
 
 private:
-  PassFPtr fn_;
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Result<Return>(Inputs &...)> fn)
+  {
+    return createImpl(name, fn);
+  }
+  template <typename Return, typename... Inputs>
+  static Pass createFn(const std::string &name,
+                       std::function<Return(Inputs...)> fn)
+  {
+    // We need to ensure that the arguments are references, but can otherwise
+    // simply wrap this in something that returns OK on the user's behalf.
+    static_assert((std::is_reference_v<Inputs> && ...),
+                  "all arguments must be references");
+    if constexpr (std::is_same_v<Return, void>) {
+      std::function<Result<OK>(Inputs...)> wrap(
+          [fn](Inputs... inputs) -> Result<OK> {
+            fn(inputs...);
+            return OK();
+          });
+      return createImpl(name, wrap);
+    } else {
+      std::function<Result<Return>(Inputs...)> wrap(
+          [fn](Inputs... inputs) -> Result<Return> { return fn(inputs...); });
+      return createImpl(name, wrap);
+    }
+  }
+  template <typename Return, typename... Inputs>
+  static Pass createImpl(const std::string &name,
+                         std::function<Result<Return>(Inputs &...)> fn)
+  {
+    std::vector<int> inputs = { PassContext::TypeId<Inputs>::type_id()... };
+    std::vector<int> outputs;
+    if constexpr (!std::is_same_v<OK, Return>) {
+      // Note that the Return value may be a reference, so we resolve the
+      // `type_id` of the underlying type. The Input types are not references,
+      // because of the way the template is constructed.
+      outputs = { PassContext::TypeId<std::decay_t<Return>>::type_id() };
+    }
 
-public:
-  std::string name;
+    return Pass(
+        name,
+        std::move(inputs),
+        std::move(outputs),
+        [fn](PassContext &ctx) -> Result<> {
+          // Extract a tuple of all input parameters, apply the function
+          // and set the result.
+          auto args = std::make_tuple(std::ref(ctx.get<Inputs>())...);
+          auto result = std::apply(fn, args);
+          if (!result) {
+            return result.takeError();
+          }
+          if constexpr (std::is_same_v<Return, OK>) {
+            return OK();
+          } else {
+            static_assert(
+                std::is_base_of_v<PassContext::State, std::decay_t<Return>>,
+                "return value must be derived from State<...>");
+            if constexpr (std::is_reference_v<Return>) {
+              // Add the external reference.
+              ctx.put(*result);
+            } else {
+              // Just move the result.
+              ctx.put(std::move(*result));
+            }
+            return OK();
+          }
+        });
+  }
+
+  // N.B. external construction is done via `create`.
+  using Fn = std::function<Result<>(PassContext &ctx)>;
+  Pass(std::string name,
+       std::vector<int> &&inputs,
+       std::vector<int> &&outputs,
+       Fn fn)
+      : name_(std::move(name)),
+        inputs_(std::move(inputs)),
+        outputs_(std::move(outputs)),
+        fn_(std::move(fn))
+  {
+  }
+
+  std::string name_;
+  std::vector<int> inputs_;
+  std::vector<int> outputs_;
+  Fn fn_;
 };
 
+// PassManager holds passes in sequence.
+//
+// The PassManager is responsible for verifying the inputs and outputs of the
+// passes ahead of time, so that the produced graph is sound.
 class PassManager {
 public:
-  PassManager(std::ostream &out = std::cerr) : out_(out) {};
+  // registers the given pass, adding to the current list.
+  PassManager &add(Pass &&pass);
 
-  void AddPass(Pass p);
-  [[nodiscard]] int Run(PassContext &ctx);
+  // registers all the given passes.
+  PassManager &add(std::vector<Pass> &&passes);
+
+  // foreach executes the given function on every pass.
+  //
+  // The first pass to encounter errors will result in execution stopping. All
+  // warnings from all passes are collected.
+  Result<> foreach(std::function<Result<>(const Pass &)> fn);
+
+  // put adds a pass which always pushes the given static data.
+  template <typename T>
+  PassManager &put(T &t)
+  {
+    return add(Pass::create(T::type_name(), [&t]() -> T & { return t; }));
+  }
+
+  // run all registered passes.
+  //
+  // This is a canonical example of how to use `foreach` above, in many other
+  // cases a direct call will be suitable.
+  Result<PassContext> run()
+  {
+    PassContext ctx;
+    auto err = foreach([&](auto &pass) -> Result<> {
+      if (!ctx.ok())
+        return OK(); // Skip the pass.
path: src/ast/pass_manager.h
user: login: amscanne
body: Fixed. I added a richer explanation to the actual method documentation, and just a note on the if here.
created_at: 2025-03-13T18:16:23Z
author_association: CONTRIBUTOR

PR comments:
user: login: amscanne
created_at: 2025-03-11T21:50:24Z
body: I have a patch pending in the stack which will split out `ASTContext` from the `PassManager`, because lots of the later passes do not depend on the AST. Will be updated shortly, but it should simplify this patch slightly.
author_association: CONTRIBUTOR
user: login: danobi
created_at: 2025-03-12T17:30:01Z
body: Quick thoughts before I disappear for the day:
* The use of template metaprogramming worries me in terms of us lay persons understanding and modifying it. However, it seems to be contained enough where the big hammer of a rewrite isn't out of the question (assuming you are hit by a bus or something). I'm sure I can fully understand it given enough time.
* It's actually not super crazy now that I'm looking at it. Seems like the trickiest bits are the parts dealing with lambdas. Which I can mostly nod through
author_association: MEMBER
user: login: jordalgo
created_at: 2025-03-12T21:58:45Z
body: > The use of template metaprogramming worries me in terms of us lay persons understanding and modifying it.

I'm also scared of too much meta programming. However, and maybe because of the in-person review @amscanne provided at Scale, I'm not as worried here. That being said, I think it all comes down to when things don't compile and I get cryptic messages that I can't resolve easily - then in a fit of rage I may curse all this stuff and attempt to rip it out. As @danobi suggested - how much would we lose by having a static list of passes and how much less meta-programming code would there be?

All that aside, these changes look pretty good to me and aside from some nits above, this is pretty much ready to go.
author_association: CONTRIBUTOR
user: login: amscanne
created_at: 2025-03-13T08:22:58Z
body: I don't think there's really any meta-programming here. It's mostly a few partial specializations and the one `TypeId` trick.

I'm not really sure that an enum is any cleaner. Do you mean something like this?

```
enum StateType {
   BpfTrace,
   AstContext,
   FooBar,
   ...
 }
 
 class PassContext {
    ...
    template <typename T>
    T& get(StateType id) {
       return *static_cast<T*>(...);
    }
 };
 ```
 
 This seems to have a lot of consistency issues and redundancy issues, e.g. you need the type and ID to correspond (`ctx.get<BPFtrace>(StateType::BpfTrace)`).
 
 You could make it so the context just have a bunch of objects directly, e.g.
 
 ```
 class PassContext {
    BPFtrace& getBpftrace() {
       return bpftrace_wrapper_.get();
    }
    std::optional<std::reference_wrapper<BPFTrace>> bpftrace_wrapper_;
  };
  ```
  
But then you just have an everything-is-about-everything class again. Really the point of this mechanism is to allow for loose couplings, e.g. you get to change some intermediate data w/o having to worry about a bunch of plumbing internals.
 
 I'm open to any other way, I'd just want a little more clarity on what you mean by a static list of passes.
author_association: CONTRIBUTOR
user: login: jordalgo
created_at: 2025-03-13T11:36:19Z
body: > I don't think there's really any meta-programming here. 

C'mon now. There is even a named technique called Template Metaprogramming, which is what is used here. You also work at Meta so it's double!

> I'm not really sure that an enum is any cleaner.

It might not be. I haven't thought about it too deeply but the thought wasn't about it being cleaner necessarily just less magical (more explicit) for folks not as familiar with the techniques being used here. Like I said, now that we've gone over it a bit in person and I reviewed this code I'm fine with things as is. 

I know @danobi said he was going to be busy this week but I'll still give him half a day (I'm so generous) to review before merging.
author_association: CONTRIBUTOR
user: login: amscanne
created_at: 2025-03-13T15:56:33Z
body: > C'mon now. There is even a named technique called Template Metaprogramming, which is what is used here. You also work at Meta so it's double!

Touché. I just mean that I don't think of it as meta programming because that's more computation or code generation through templates. I think I managed to avoid a lot of that.

> It might not be. I haven't thought about it too deeply but the thought wasn't about it being cleaner necessarily just less magical (more explicit) for folks not as familiar with the techniques being used here. Like I said, now that we've gone over it a bit in person and I reviewed this code I'm fine with things as is.
> 
> I know @danobi said he was going to be busy this week but I'll still give him half a day (I'm so generous) to review before merging.

I don't want to foist anything, so I'm open to exploring other avenues. I recognize that templates can be used in good taste or bad. I tend to think it's best when it's compartmentalized to some core functionality, but then it enables a broad set of patterns (which aren't directly loaded with magic). This is what I tried to do here. I did a pass at this last night to simplify (dropped all the `TypeName` stuff, for example), but ultimately felt that could hurt debuggability.

In terms of how this all manifests, I think it's best to look ahead at the [codegen passes](https://github.com/bpftrace/bpftrace/pull/3885/files#diff-28926c090177e6efbfb5a7449833e0e863bb226822abb513b7dfc6c388837d87R4524) to evaluate if it's worth it.
author_association: CONTRIBUTOR
user: login: viktormalik
created_at: 2025-03-13T17:01:23Z
body: > In terms of how this all manifests, I think it's best to look ahead at the [codegen passes](https://github.com/bpftrace/bpftrace/pull/3885/files#diff-28926c090177e6efbfb5a7449833e0e863bb226822abb513b7dfc6c388837d87R4524) to evaluate if it's worth it.

I like how this allows to split these large passes we have into a number of smaller, interconnected passes while also being able to check that the dependencies among them are contained.

I wish it could be done without the template magic under the hood but I'm ok to trade off the complexity in the core part for simplicity to use, which seems to be the case here.
author_association: CONTRIBUTOR