[willow]: add ConstantPool (#5)

4 files changed, 161 insertions, 2 deletions

diff --git a/willow/include/willow/IR/Constant.h b/willow/include/willow/IR/Constant.h
new file mode 100644
index 0000000..d93d65d
--- /dev/null
+++ b/willow/include/willow/IR/Constant.h
@@ -0,0 +1,92 @@
+#ifndef WILLOW_INCLUDE_IR_CONSTANT_H
+#define WILLOW_INCLUDE_IR_CONSTANT_H
+
+#include <willow/IR/Value.h>
+
+namespace willow {
+
+enum class ConstantKind {
+  Int,    //< Integer value with known bits
+  Undef,  //< Known undef
+  Poison, //< Known poison
+};
+
+class Constant : public Value {
+  ConstantKind kind;
+
+protected:
+  Constant(Type type, ConstantKind k)
+      : Value(ValueKind::Constant, type), kind(k) {}
+
+public:
+  Constant(const Constant &) = delete;
+  Constant &operator=(const Constant &) = delete;
+
+  bool isUndef() const { return kind == ConstantKind::Undef; }
+  bool isPoison() const { return kind == ConstantKind::Poison; }
+  bool isInt() const { return kind == ConstantKind::Int; }
+};
+
+/// A constant integer with a known value
+class ConstantInt final : public Constant {
+  // TODO: bigint
+  uint64_t bits = 0;
+
+public:
+  explicit ConstantInt(Type ty, uint64_t val)
+      : Constant(ty, ConstantKind::Int) {
+    const size_t w = ty.getNumBits();
+    assert(w <= 64 && "Come back when im less lazy");
+    // assert(!(val >> w) && "Truncated constant");
+
+    if (w == 64)
+      bits = val;
+    else {
+      const uint64_t m = (uint64_t{1} << w) - 1;
+      bits = val & m;
+    }
+  }
+
+  int64_t getSExtVal() const {
+    const size_t w = getType().getNumBits();
+
+    if (w == 64)
+      return static_cast<int64_t>(bits);
+
+    unsigned sh = 64 - w;
+    return (static_cast<int64_t>(bits << sh)) >> sh;
+  }
+
+  uint64_t getZExtVal() const { return bits; }
+
+  struct Key {
+    TypeImpl *ty;
+    uint64_t bits; // canonicalized value
+    bool operator==(const ConstantInt::Key &) const = default;
+  };
+
+  struct Hash {
+    size_t operator()(const ConstantInt::Key &k) const noexcept {
+      auto h1 = std::hash<TypeImpl *>{}(k.ty);
+      auto h2 = std::hash<uint64_t>{}(k.bits);
+
+      return h1 ^ (h2 + 0x9e3779b97f4a7c15ULL + (h1 << 6) + (h1 >> 2));
+    }
+  };
+};
+
+/// Represents a constant value that is entirely undef.
+class UndefVal final : public Constant {
+public:
+  explicit UndefVal(Type ty) : Constant(ty, ConstantKind::Undef) {}
+};
+
+/// Represents a constant value that is entirely poison.
+class PoisonVal final : public Constant {
+public:
+  explicit PoisonVal(Type ty) : Constant(ty, ConstantKind::Poison) {}
+};
+
+} // namespace willow
+
+#endif // WILLOW_INCLUDE_IR_CONSTANT_H
diff --git a/willow/include/willow/IR/ConstantPool.h b/willow/include/willow/IR/ConstantPool.h
new file mode 100644
index 0000000..2524c70
--- /dev/null
+++ b/willow/include/willow/IR/ConstantPool.h
@@ -0,0 +1,65 @@
+#ifndef WILLOW_INCLUDE_IR_CONSTANT_POOL_H
+#define WILLOW_INCLUDE_IR_CONSTANT_POOL_H
+
+#include <memory>
+#include <willow/IR/Constant.h>
+
+namespace willow {
+
+class ConstantPool {
+public:
+  ConstantPool() = default;
+  ConstantPool(const ConstantPool &) = delete;
+  ConstantPool &operator=(const ConstantPool &) = delete;
+
+  /// Get a uniqued integer constant with the specified type and value.
+  ConstantInt *getInt(Type ty, uint64_t val);
+  /// Get a uniqued undef value for the specified type \p ty.
+  UndefVal *getUndefVal(Type ty);
+  /// Get a uniqued poison value for the specified type \p ty.
+  PoisonVal *getPoisonVal(Type ty);
+
+private:
+  std::mutex int_mutex;
+  std::unordered_map<ConstantInt::Key, std::unique_ptr<ConstantInt>,
+                     ConstantInt::Hash>
+      icache;
+
+  std::mutex undef_mutex;
+  std::unordered_map<TypeImpl *, std::unique_ptr<UndefVal>> ucache;
+
+  std::mutex poison_mutex;
+  std::unordered_map<TypeImpl *, std::unique_ptr<PoisonVal>> pcache;
+};
+
+ConstantInt *ConstantPool::getInt(Type ty, uint64_t val) {
+  assert(ty.isInt() && "Expected integer type");
+  ConstantInt::Key &&k{ty.getImpl(), ty.getNumBits()};
+  std::lock_guard<std::mutex> lock(int_mutex);
+
+  auto [it, _] = icache.try_emplace(k, std::make_unique<ConstantInt>(ty, val));
+
+  return it->second.get();
+}
+
+UndefVal *ConstantPool::getUndefVal(Type ty) {
+  std::lock_guard<std::mutex> lock(undef_mutex);
+
+  auto [it, _] =
+      ucache.try_emplace(ty.getImpl(), std::make_unique<UndefVal>(ty));
+
+  return it->second.get();
+}
+
+PoisonVal *ConstantPool::getPoisonVal(Type ty) {
+  std::lock_guard<std::mutex> lock(poison_mutex);
+
+  auto [it, _] =
+      pcache.try_emplace(ty.getImpl(), std::make_unique<PoisonVal>(ty));
+
+  return it->second.get();
+}
+
+}; // namespace willow
+
+#endif // WILLOW_INCLUDE_IR_CONSTANT_POOL_H
diff --git a/willow/include/willow/IR/TypeContext.h b/willow/include/willow/IR/TypeContext.h
index 12aaa02..3d414b1 100644
--- a/willow/include/willow/IR/TypeContext.h
+++ b/willow/include/willow/IR/TypeContext.h
@@ -38,7 +38,8 @@ private:
 };
 
 Type TypeContext::getIntType(std::size_t width) {
-  auto [it, _] = icache.try_emplace(IntTypeImpl::Key{width});
+  auto [it, _] = icache.try_emplace(IntTypeImpl::Key{width},
+                                    std::make_unique<IntTypeImpl>(width));
 
   return Type(it->second.get());
 }
diff --git a/willow/include/willow/IR/Value.h b/willow/include/willow/IR/Value.h
index 138b895..7c93728 100644
--- a/willow/include/willow/IR/Value.h
+++ b/willow/include/willow/IR/Value.h
@@ -15,7 +15,7 @@ class Instruction;
 enum class ValueKind {
   Instruction, ///< the identified result of an instruction
   Parameter,   ///< the named parameter to a function
-  Literal,     ///< an anonymous typed literal
+  Constant,    ///< an anonymous typed constant
 };
 
 /// An SSA value that may be used.
@@ -40,6 +40,7 @@ public:
   void setName(std::string name) { this->name = std::move(name); }
 
   Type getType() const { return type; }
+  // TODO: less confusing name
   ValueKind getValueType() const { return value_type; }
 
   bool isVoid() const { return type.isVoid(); }