refactor xir constant

include/luisa/xir/constant.h

@@ -1,22 +1,42 @@
 #pragma once
 
-#include <luisa/ast/constant_data.h>
 #include <luisa/xir/value.h>
 
 namespace luisa::compute::xir {
 
 class LC_XIR_API Constant : public Value {
 
 private:
-    ConstantData _data;
+    union {
+        std::byte _small[sizeof(void *)] = {};
+        void *_large;
+    };
+
+    [[nodiscard]] bool _is_small() const noexcept;
 
 public:
-    explicit Constant(Pool *pool, ConstantData data = {}, const Name *name = nullptr) noexcept;
-    void set_data(ConstantData data) noexcept;
-    [[nodiscard]] auto data() const noexcept { return _data; }
+    explicit Constant(Pool *pool, const Type *type,
+                      const void *data = nullptr,
+                      const Name *name = nullptr) noexcept;
+    ~Constant() noexcept override;
+    void set_type(const Type *type) noexcept override;
+    void set_data(const void *data) noexcept;
+    [[nodiscard]] void *data() noexcept;
+    [[nodiscard]] const void *data() const noexcept;
     [[nodiscard]] DerivedValueTag derived_value_tag() const noexcept final {
         return DerivedValueTag::CONSTANT;
     }
+
+    template<typename T>
+    [[nodiscard]] T &as() noexcept {
+        assert(type()->size() == sizeof(T) && "Type size mismatch.");
+        return *reinterpret_cast<T *>(data());
+    }
+
+    template<typename T>
+    [[nodiscard]] const T &as() const noexcept {
+        return const_cast<Constant *>(this)->as<T>();
+    }
 };
 
 }// namespace luisa::compute::xir

include/luisa/xir/value.h

@@ -31,8 +31,8 @@ class LC_XIR_API Value : public PooledObject {
     explicit Value(Pool *pool, const Type *type = nullptr, const Name *name = nullptr) noexcept;
     [[nodiscard]] virtual DerivedValueTag derived_value_tag() const noexcept = 0;
 
-    void set_type(const Type *type) noexcept { _type = type; }
-    void set_name(const Name *name) noexcept { _name = name; }
+    virtual void set_type(const Type *type) noexcept { _type = type; }
+    virtual void set_name(const Name *name) noexcept { _name = name; }
 
     [[nodiscard]] auto type() const noexcept { return _type; }
     [[nodiscard]] auto name() const noexcept { return _name; }

src/xir/constant.cpp

@@ -1,17 +1,47 @@
+#include <luisa/ast/type_registry.h>
 #include <luisa/core/logging.h>
 #include <luisa/xir/constant.h>
 
 namespace luisa::compute::xir {
 
-Constant::Constant(Pool *pool, ConstantData data, const Name *name) noexcept
-    : Value{pool, data.type(), name} { set_data(data); }
+Constant::Constant(Pool *pool, const Type *type, const void *data, const Name *name) noexcept
+    : Value{pool, type, name} {
+    LUISA_ASSERT(type != nullptr, "Constant type must be specified.");
+    if (!_is_small()) {
+        _large = luisa::allocate_with_allocator<std::byte>(type->size());
+    }
+    set_data(data);
+}
+
+Constant::~Constant() noexcept {
+    if (!_is_small()) {
+        luisa::deallocate_with_allocator(static_cast<std::byte *>(_large));
+    }
+}
+
+bool Constant::_is_small() const noexcept {
+    return type()->size() <= sizeof(void *);
+}
+
+void Constant::set_data(const void *data) noexcept {
+    if (data == nullptr) {
+        memset(this->data(), 0, type()->size());
+    } else {
+        memmove(this->data(), data, type()->size());
+    }
+}
+
+void Constant::set_type(const Type *type) noexcept {
+    LUISA_ERROR_WITH_LOCATION("Constant type cannot be changed.");
+}
+
+void *Constant::data() noexcept {
+    if (type() == nullptr) { return nullptr; }
+    return type()->size() <= sizeof(void *) ? _small : _large;
+}
 
-void Constant::set_data(ConstantData data) noexcept {
-    LUISA_DEBUG_ASSERT(!data || data.type() == type(),
-                       "Constant data type mismatch: {} vs {}",
-                       data.type()->description(),
-                       type()->description());
-    _data = data;
+const void *Constant::data() const noexcept {
+    return const_cast<Constant *>(this)->data();
 }
 
 }// namespace luisa::compute::xir