Expand the tower component.

include/vast/Tower/Handle.hpp

@@ -0,0 +1,22 @@
+// Copyright (c) 2022-present, Trail of Bits, Inc.
+
+#pragma once
+
+#include "vast/Util/Common.hpp"
+
+VAST_RELAX_WARNINGS
+VAST_UNRELAX_WARNINGS
+
+namespace vast::tw {
+    using conversion_path_t             = std::vector< std::string >;
+    using conversion_path_fingerprint_t = std::string;
+
+    using handle_id_t = std::size_t;
+
+    struct handle_t
+    {
+        handle_id_t id;
+        vast_module mod;
+    };
+
+} // namespace vast::tw

include/vast/Tower/Link.hpp

@@ -0,0 +1,118 @@
+// Copyright (c) 2024-present, Trail of Bits, Inc.
+
+#pragma once
+
+#include "vast/Util/Common.hpp"
+
+#include "vast/Tower/Handle.hpp"
+#include "vast/Tower/LocationInfo.hpp"
+
+VAST_RELAX_WARNINGS
+VAST_UNRELAX_WARNINGS
+
+#include <memory>
+#include <vector>
+
+namespace vast::tw {
+
+    using operations = std::vector< operation >;
+
+    // Maybe we want to abstract this as an interface - then it can be lazy as well, for example
+    // we could pass in only some step chain and location info and it will get computed.
+    // Since if we have `A -> B` we can always construct `B -> A` it may make sense to simply only
+    // export bidirectional mapping?
+    using op_mapping = llvm::DenseMap< operation, operations >;
+
+    // Generic interface to generalize the connection between any two modules.
+    // There are no performance guarantees in general, but there should be implementations
+    // available that try to be as performant as possible.
+    struct link_interface
+    {
+        virtual ~link_interface() = default;
+
+        // These are not forced as `const` to allow runtime caching.
+        virtual operations children(operation)  = 0;
+        virtual operations children(operations) = 0;
+
+        virtual operations parents(operation)  = 0;
+        virtual operations parents(operations) = 0;
+
+        virtual op_mapping parents_to_children() = 0;
+        virtual op_mapping children_to_parents() = 0;
+
+        virtual handle_t parent() const = 0;
+        virtual handle_t child() const = 0;
+    };
+
+    namespace views {
+        static inline auto parents_to_children = [](const auto &link) {
+            return link->parents_to_children();
+        };
+
+        static inline auto children_to_parents = [](const auto &link) {
+            return link->children_to_parents();
+        };
+
+    } // namespace views
+
+    // Represent application of some passes. Invariant is that
+    // `parent -> child` are tied by the `location_info`.
+    // TODO: How to enforce this - private ctor and provide a builder interface on the side
+    //       that is a friend and allowed to create these?
+    struct conversion_step : link_interface {
+      protected:
+        handle_t _parent;
+        handle_t _child;
+        location_info_t &_location_info;
+
+      public:
+        explicit conversion_step(handle_t parent, handle_t child, location_info_t &location_info)
+            : _parent(parent), _child(child), _location_info(location_info)
+        {}
+
+        operations children(operation) override;
+        operations children(operations) override;
+
+        operations parents(operation) override;
+        operations parents(operations) override;
+
+        op_mapping parents_to_children() override;
+        op_mapping children_to_parents() override;
+
+        handle_t parent() const override;
+        handle_t child() const override;
+    };
+
+    using link_ptr = std::unique_ptr< link_interface >;
+    using link_vector = std::vector< link_ptr >;
+
+    using conversion_steps = std::vector< conversion_step >;
+
+    // `A -> ... -> E` - each middle link is kept and there is pre-computed
+    // mapping for `A <-> E` transition to make it more performant.
+    struct fat_link : link_interface
+    {
+      protected:
+        link_vector _links;
+
+        op_mapping _to_children;
+        op_mapping _to_parents;
+
+      public:
+        explicit fat_link(link_vector links);
+        fat_link() = delete;
+
+        operations children(operation) override;
+        operations children(operations) override;
+
+        operations parents(operation) override;
+        operations parents(operations) override;
+
+        op_mapping parents_to_children() override;
+        op_mapping children_to_parents() override;
+
+        handle_t child() const override;
+        handle_t parent() const override;
+    };
+
+} // namespace vast::tw

include/vast/Tower/LocationInfo.hpp

@@ -0,0 +1,68 @@
+// Copyright (c) 2022-present, Trail of Bits, Inc.
+
+#pragma once
+
+#include "vast/Util/Common.hpp"
+
+VAST_RELAX_WARNINGS
+#include <mlir/Pass/Pass.h>
+#include <mlir/Pass/PassManager.h>
+VAST_UNRELAX_WARNINGS
+
+#include "vast/Tower/Handle.hpp"
+
+namespace vast::tw {
+
+    // Is allowed to have state?
+    struct location_info_t
+    {
+      private:
+        // Encoded as `mlir::FusedLocation(original, mlir::OpaqueLocation(pointer_to_self))`
+        using raw_loc_t = mlir::FusedLoc;
+
+        static raw_loc_t raw_loc(operation op) {
+            auto raw = mlir::dyn_cast< raw_loc_t >(op->getLoc());
+            VAST_CHECK(raw, "{0} with loc: {1}", *op, op->getLoc());
+            return raw;
+        }
+
+        template< std::size_t idx >
+            requires (idx < 2)
+        static loc_t get(raw_loc_t raw) {
+            auto locs = raw.getLocations();
+            VAST_ASSERT(locs.size() == 2);
+            return locs[idx];
+        }
+
+        static auto parse(operation op) { return std::make_tuple(prev(op), self(op)); }
+
+        // TODO: These are strictly not needed in this form, but help initial
+        //       debugging a lot.
+        std::string fingerprint(const conversion_path_t &);
+        loc_t mk_unique_loc(const conversion_path_t &, operation);
+        loc_t mk_linked_loc(loc_t self, loc_t prev);
+
+      public:
+        // For the given operation return location to be used in this module.
+        loc_t get_as_child(const conversion_path_t &, operation op);
+        loc_t get_root(operation op);
+
+        static loc_t self(raw_loc_t raw) { return get< 1 >(raw); }
+
+        static loc_t prev(raw_loc_t raw) { return get< 0 >(raw); }
+
+        static loc_t self(operation op) { return self(raw_loc(op)); }
+
+        static loc_t prev(operation op) { return prev(raw_loc(op)); }
+
+        static bool are_tied(operation parent, operation child);
+    };
+
+    // Since we are going to tie together arbitrary modules, it makes sense to make them
+    // have locations in the same shape - therefore root shouldn't be an excuse. It will
+    // however require slightly different handling, so we are exposing a hook for that.
+    void mk_root(location_info_t &, operation);
+
+    void transform_locations(location_info_t &, const conversion_path_t &, operation);
+
+} // namespace vast::tw

include/vast/Tower/Storage.hpp

@@ -0,0 +1,52 @@
+// Copyright (c) 2022-present, Trail of Bits, Inc.
+
+#pragma once
+
+#include "vast/Util/Common.hpp"
+
+VAST_RELAX_WARNINGS
+#include <mlir/Pass/Pass.h>
+#include <mlir/Pass/PassManager.h>
+VAST_UNRELAX_WARNINGS
+
+#include "vast/Tower/Handle.hpp"
+
+#include <algorithm>
+#include <numeric>
+
+namespace vast::tw {
+
+    struct module_storage
+    {
+        // TODO: API-wise, we probably want to accept any type that is `mlir::OwningOpRef< T >`?
+        handle_t store(const conversion_path_t &path, owning_module_ref mod) {
+            auto id      = allocate_id(path);
+            auto [it, _] = storage.insert({id, std::move(mod)});
+            return { id, it->second.get() };
+        }
+
+        void remove(handle_t) { VAST_UNIMPLEMENTED; }
+
+      private:
+        conversion_path_fingerprint_t fingerprint(const conversion_path_t &path) const {
+            return std::accumulate(path.begin(), path.end(), std::string{});
+        }
+
+        handle_id_t allocate_id(const conversion_path_t &path) {
+            return allocate_id(fingerprint(path));
+        }
+
+        handle_id_t allocate_id(const conversion_path_fingerprint_t &fp) {
+            // Later here we want to return the cached module?
+            VAST_CHECK(!conversion_tree.count(fp), "For now cannot do caching!");
+            auto id = next_id++;
+            conversion_tree.emplace(fp, id);
+            return id;
+        }
+
+        std::size_t next_id = 0;
+        llvm::DenseMap< handle_id_t, owning_module_ref > storage;
+        // TODO: This is just a prototyping shortcut, we may want something smarter here.
+        std::unordered_map< conversion_path_fingerprint_t, handle_id_t > conversion_tree;
+    };
+} // namespace vast::tw

include/vast/Tower/Tower.hpp

@@ -5,75 +5,38 @@
 #include "vast/Util/Common.hpp"
 
 VAST_RELAX_WARNINGS
-#include <mlir/Pass/PassManager.h>
 #include <mlir/Pass/Pass.h>
+#include <mlir/Pass/PassManager.h>
 VAST_UNRELAX_WARNINGS
 
-namespace vast::tw {
-
-    struct default_loc_rewriter_t
-    {
-        static auto insert(mlir::Operation *op) -> void;
-        static auto remove(mlir::Operation *op) -> void;
-        static auto prev(mlir::Operation *op) -> mlir::Operation *;
-    };
+#include "vast/Tower/Handle.hpp"
+#include "vast/Tower/Link.hpp"
+#include "vast/Tower/LocationInfo.hpp"
+#include "vast/Tower/Storage.hpp"
 
-    using pass_ptr_t = std::unique_ptr< mlir::Pass >;
+namespace vast::tw {
 
-    template< typename loc_rewriter_t >
     struct tower
     {
-        using loc_rewriter = loc_rewriter_t;
-
-        struct handle_t
-        {
-            std::size_t id;
-            vast_module mod;
-        };
-
-        static auto get(mcontext_t &ctx, owning_module_ref mod)
-            -> std::tuple< tower, handle_t > {
-            tower t(ctx, std::move(mod));
-            handle_t h{ .id = 0, .mod = t._modules[0].get() };
-            return { std::move(t), h };
-        }
-
-        auto apply(handle_t handle, mlir::PassManager &pm) -> handle_t {
-            handle.mod.walk(loc_rewriter::insert);
-
-            _modules.emplace_back(mlir::cast< vast_module >(handle.mod->clone()));
-
-            auto id  = _modules.size() - 1;
-            auto mod = _modules.back().get();
-
-            if (mlir::failed(pm.run(mod))) {
-                VAST_FATAL("some pass in apply() failed");
-            }
-
-            handle.mod.walk(loc_rewriter::remove);
-
-            return { id, mod };
-        }
-
-        auto apply(handle_t handle, pass_ptr_t pass) -> handle_t {
-            mlir::PassManager pm(_ctx);
-            pm.addPass(std::move(pass));
-            return apply(handle, pm);
+      private:
+        [[maybe_unused]] mcontext_t &mctx;
+        module_storage storage;
+        handle_t top_handle;
+
+      public:
+        tower(mcontext_t &mctx, location_info_t &li, owning_module_ref root) : mctx(mctx) {
+            mk_root(li, root->getOperation());
+            top_handle = storage.store(root_conversion(), std::move(root));
         }
 
-        auto top() -> handle_t { return { _modules.size(), _modules.back().get() }; }
-
       private:
-        using module_storage_t = llvm::SmallVector< owning_module_ref, 2 >;
+        // TODO: Move somewhere else.
+        static conversion_path_t root_conversion() { return {}; }
 
-        mcontext_t *_ctx;
-        module_storage_t _modules;
+      public:
+        handle_t top() const { return top_handle; }
 
-        tower(mcontext_t &ctx, owning_module_ref mod) : _ctx(&ctx) {
-            _modules.emplace_back(std::move(mod));
-        }
+        link_ptr apply(handle_t, location_info_t &, mlir::PassManager &);
     };
 
-    using default_tower = tower< default_loc_rewriter_t >;
-
 } // namespace vast::tw