fix type conversions

EnzymeAD · Aug 12, 2023 · d6c628a · d6c628a
1 parent a7db7a6
commit d6c628a
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 12 deletions.
diff --git a/enzyme/Enzyme/TypeAnalysis/Trie.h b/enzyme/Enzyme/TypeAnalysis/Trie.h
@@ -7,6 +7,7 @@
 #include <map>
 #include <limits>
 #include <stdexcept>
+#include <cassert>
 
 #define private public
 
@@ -184,7 +185,7 @@ class trie_const_iterator {
 
   explicit trie_const_iterator(const trie<Key, T, Allocator> *trie, const trie_node<Key, Index, value_type> *n) : _trie(trie), _node{n} {}
 
-  explicit trie_const_iterator(const trie_iterator<Key, Index, T, Allocator> iter) : _trie(iter._trie), _node{iter._node} {}
+  trie_const_iterator(const trie_iterator<Key, Index, T, Allocator> iter) : _trie(iter._trie), _node{iter._node} {}
 
   trie_const_iterator(const trie_const_iterator &other) : _trie(other._trie), _node(other._node) {}
 

diff --git a/enzyme/Enzyme/TypeAnalysis/TypeTree.h b/enzyme/Enzyme/TypeAnalysis/TypeTree.h
@@ -41,6 +41,7 @@
 #include "../Utils.h"
 #include "BaseType.h"
 #include "ConcreteType.h"
+#include "Trie.h"
 
 /// Maximum offset for type trees to keep
 extern "C" {
@@ -64,15 +65,13 @@ static inline std::string to_string(const std::vector<int> x) {
 class TypeTree;
 
 typedef std::shared_ptr<const TypeTree> TypeResult;
-typedef std::map<const std::vector<int>, ConcreteType> ConcreteTypeMapType;
-typedef std::map<const std::vector<int>, const TypeResult> TypeTreeMapType;
 
 /// Class representing the underlying types of values as
 /// sequences of offsets to a ConcreteType
 class TypeTree : public std::enable_shared_from_this<TypeTree> {
 private:
   // mapping of known indices to type if one exists
-  ConcreteTypeMapType mapping;
+  trie<std::vector<int>, ConcreteType> mapping;
   std::vector<int> minIndices;
 
 public:
@@ -84,7 +83,7 @@ class TypeTree : public std::enable_shared_from_this<TypeTree> {
   }
 
   /// Utility helper to lookup the mapping
-  const ConcreteTypeMapType &getMapping() const { return mapping; }
+  const trie<std::vector<int>, ConcreteType> &getMapping() const { return mapping; }
 
   /// Lookup the underlying ConcreteType at a given offset sequence
   /// or Unknown if none exists
@@ -888,7 +887,7 @@ class TypeTree : public std::enable_shared_from_this<TypeTree> {
     minIndices = RHS.minIndices;
     mapping.clear();
     for (const auto &elems : RHS.mapping) {
-      mapping.emplace(elems);
+      mapping.insert(elems);
     }
     return true;
   }
@@ -1121,20 +1120,20 @@ class TypeTree : public std::enable_shared_from_this<TypeTree> {
     }
 
     // mapings just on the right
-    for (auto &pair : RHS.mapping) {
+    for (auto&& [key, ct] : RHS.mapping) {
       // TODO propagate non-first level operands:
       // Special handling is necessary here because a pointer to an int
       // binop with something should not apply the binop rules to the
       // underlying data but instead a different rule
-      if (pair.first.size() > 0) {
+      if (key.size() > 0) {
         continue;
       }
-
-      if (mapping.find(pair.first) == RHS.mapping.end()) {
+                
+      if (RHS.mapping.end() == mapping.find(key)) {
         ConcreteType CT = BaseType::Unknown;
-        changed |= CT.binopIn(pair.second, Op);
+        changed |= CT.binopIn(ct, Op);
         if (CT != BaseType::Unknown) {
-          mapping.insert(std::make_pair(pair.first, CT));
+          mapping.insert(std::make_pair(key, CT));
         }
       }
     }