1 files changed, 6 insertions, 2 deletions
diff --git a/include/CompilerKit/CodeGenerator.h b/include/CompilerKit/CodeGenerator.h
index 8119c58..721743d 100644
--- a/include/CompilerKit/CodeGenerator.h
+++ b/include/CompilerKit/CodeGenerator.h
@@ -8,6 +8,10 @@
 
 #include <CompilerKit/Detail/Config.h>
 #include <CompilerKit/Macros.h>
+
+#include <CompilerKit/ErrorOr.h>
+#include <CompilerKit/Ref.h>
+
 #include <cstring>
 
 #define CK_ASSEMBLY_INTERFACE : public ::CompilerKit::IAssembly
@@ -42,8 +46,8 @@ class AssemblyFactory final {
 
   Int32 Compile(STLString sourceFile, const Int32& arch) noexcept;
 
-  void       Mount(IAssembly* mountPtr) noexcept;
-  IAssembly* Unmount() noexcept;
+  void       Mount(WeakRef<IAssembly> mountPtr) noexcept;
+  WeakRef<IAssembly> Unmount() noexcept;
 
  private:
   IAssembly* fMounted{nullptr};