[CIR-link] Refactor nest/updateState (#35)

* [CIR-link] Refactor nest/updateState

* Formatting

Author: hbrodin

mlir/include/mlir/Linker/LLVMLinkerMixin.h

@@ -191,13 +191,13 @@ class LLVMLinkerMixin {
              isAvailableExternallyLinkage(srcLinkage));
   }
 
-  LogicalResult verifyLinkageCompatibility(Conflict pair) const{
+  LogicalResult verifyLinkageCompatibility(Conflict pair) const {
     const DerivedLinkerInterface &derived = getDerived();
     assert(derived.canBeLinked(pair.src) && "expected linkable operation");
     assert(derived.canBeLinked(pair.dst) && "expected linkable operation");
 
     auto linkError = [&](const Twine &error) -> LogicalResult {
-        return pair.src->emitError(error) << " dst: " << pair.dst->getLoc();
+      return pair.src->emitError(error) << " dst: " << pair.dst->getLoc();
     };
 
     Linkage srcLinkage = derived.getLinkage(pair.src);
@@ -208,13 +208,14 @@ class LLVMLinkerMixin {
 
     if (isAppendingLinkage(srcLinkage) && isAppendingLinkage(dstLinkage)) {
       if (srcUnnamedAddr != dstUnnamedAddr) {
-        return linkError("Appending variables with different unnamed_addr need to be linked");
+        return linkError("Appending variables with different unnamed_addr need "
+                         "to be linked");
       }
     }
     return success();
   }
 
-  ConflictResolution getConflictResolution(Conflict pair) const  {
+  ConflictResolution getConflictResolution(Conflict pair) const {
     const DerivedLinkerInterface &derived = getDerived();
     assert(derived.canBeLinked(pair.src) && "expected linkable operation");
     assert(derived.canBeLinked(pair.dst) && "expected linkable operation");
@@ -318,7 +319,6 @@ class SymbolAttrLLVMLinkerInterface
   ConflictResolution getConflictResolution(Conflict pair) const override {
     return LinkerMixin::getConflictResolution(pair);
   }
-
 };
 
 } // namespace mlir::link

mlir/include/mlir/Linker/LinkerInterface.h

@@ -20,6 +20,7 @@
 #include "mlir/IR/IRMapping.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/Support/Error.h"
+#include <memory>
 
 namespace mlir::link {
 
@@ -35,7 +36,8 @@ enum LinkerFlags {
 
 class LinkState {
 public:
-  LinkState(ModuleOp dst) : builder(dst.getBodyRegion()) {}
+  LinkState(ModuleOp dst)
+      : mapping(std::make_shared<IRMapping>()), builder(dst.getBodyRegion()) {}
 
   Operation *clone(Operation *src);
   Operation *cloneWithoutRegions(Operation *src);
@@ -46,10 +48,12 @@ class LinkState {
 
   LinkState nest(ModuleOp submod) const;
 
-  void updateState(const LinkState &substate);
-
 private:
-  IRMapping mapping;
+  // Private constructor used by nest()
+  LinkState(ModuleOp dst, std::shared_ptr<IRMapping> mapping)
+      : mapping(mapping), builder(dst.getBodyRegion()) {}
+
+  std::shared_ptr<IRMapping> mapping;
   OpBuilder builder;
 };
 
@@ -167,7 +171,8 @@ class SymbolAttrLinkerInterface : public SymbolLinkerInterface {
   /// Resolves a conflict between an existing operation and a new one.
   LogicalResult resolveConflict(Conflict pair) override;
 
-  virtual LogicalResult resolveConflict(Conflict pair, ConflictResolution resolution);
+  virtual LogicalResult resolveConflict(Conflict pair,
+                                        ConflictResolution resolution);
 
   /// Gets the conflict resolution for a given conflict
   virtual ConflictResolution getConflictResolution(Conflict pair) const = 0;

mlir/lib/Linker/LinkerInterface.cpp

@@ -18,22 +18,23 @@ using namespace mlir::link;
 //===----------------------------------------------------------------------===//
 
 template <typename CloneFunc>
-Operation *cloneImpl(Operation *src, IRMapping &mapping, CloneFunc cloneFunc) {
-  assert(!mapping.contains(src));
+Operation *cloneImpl(Operation *src, std::shared_ptr<IRMapping> &mapping,
+                     CloneFunc cloneFunc) {
+  assert(!mapping->contains(src));
   Operation *dst = cloneFunc(src);
-  mapping.map(src, dst);
+  mapping->map(src, dst);
   return dst;
 }
 
 Operation *LinkState::clone(Operation *src) {
   return cloneImpl(src, mapping, [this](Operation *op) {
-    return builder.clone(*op, mapping);
+    return builder.clone(*op, *mapping);
   });
 }
 
 Operation *LinkState::cloneWithoutRegions(Operation *src) {
   return cloneImpl(src, mapping, [this](Operation *op) {
-    return builder.cloneWithoutRegions(*op, mapping);
+    return builder.cloneWithoutRegions(*op, *mapping);
   });
 }
 
@@ -42,20 +43,14 @@ Operation *LinkState::getDestinationOp() const {
 }
 
 Operation *LinkState::remapped(Operation *src) const {
-  return mapping.lookupOrNull(src);
+  return mapping->lookupOrNull(src);
 }
 
 LinkState LinkState::nest(ModuleOp submod) const {
   assert(submod->getParentOfType<mlir::ModuleOp>().getOperation() ==
              getDestinationOp() &&
          "Submodule should be directly nested in the current state");
-  LinkState submodState(submod);
-  submodState.mapping = mapping;
-  return submodState;
-}
-
-void LinkState::updateState(const LinkState &substate) {
-  mapping = substate.mapping;
+  return LinkState(submod, mapping);
 }
 
 //===----------------------------------------------------------------------===//