feat(core): Support deepcopy in opaque (#76)

Author: potatomashed

cpp/structure.cc

@@ -627,7 +627,7 @@ inline void StructuralEqualImpl(Object *lhs, Object *rhs, bool bind_free_vars) {
     }
     // `task.visited` was `False`
     int64_t task_index = static_cast<int64_t>(tasks.size()) - 1;
-    if (type_info->type_index == kMLCList) {
+    if (lhs->IsInstance<UListObj>()) {
       UListObj *lhs_list = reinterpret_cast<UListObj *>(lhs);
       UListObj *rhs_list = reinterpret_cast<UListObj *>(rhs);
       int64_t lhs_size = lhs_list->size();
@@ -639,7 +639,7 @@ inline void StructuralEqualImpl(Object *lhs, Object *rhs, bool bind_free_vars) {
         auto &err = tasks[task_index].err = std::make_unique<std::ostringstream>();
         (*err) << "List length mismatch: " << lhs_size << " vs " << rhs_size;
       }
-    } else if (type_info->type_index == kMLCDict) {
+    } else if (lhs->IsInstance<UDictObj>()) {
       UDictObj *lhs_dict = reinterpret_cast<UDictObj *>(lhs);
       UDictObj *rhs_dict = reinterpret_cast<UDictObj *>(rhs);
       std::vector<AnyView> not_found_lhs_keys;
@@ -892,13 +892,13 @@ inline uint64_t StructuralHashImpl(Object *obj) {
       task.index_in_result_hashes = result_hashes.size();
     }
     // `task.visited` was `False`
-    if (type_info->type_index == kMLCList) {
+    if (obj->IsInstance<UListObj>()) {
       UListObj *list = reinterpret_cast<UListObj *>(obj);
       hash_value = HashCombine(hash_value, list->size());
       for (int64_t i = list->size() - 1; i >= 0; --i) {
         Visitor::EnqueueAny(&tasks, bind_free_vars, &list->at(i));
       }
-    } else if (type_info->type_index == kMLCDict) {
+    } else if (obj->IsInstance<UDictObj>()) {
       UDictObj *dict = reinterpret_cast<UDictObj *>(obj);
       hash_value = HashCombine(hash_value, dict->size());
       struct KVPair {
@@ -1151,8 +1151,16 @@ inline Any CopyDeepImpl(AnyView source) {
     } else if (object->IsInstance<StrObj>() || object->IsInstance<ErrorObj>() || object->IsInstance<FuncObj>() ||
                object->IsInstance<TensorObj>()) {
       ret = object;
-    } else if (object->IsInstance<OpaqueObj>()) {
-      MLC_THROW(TypeError) << "Cannot copy `mlc.Opaque` of type: " << object->DynCast<OpaqueObj>()->opaque_type_name;
+    } else if (OpaqueObj *opaque = object->as<OpaqueObj>()) {
+      std::string func_name = "Opaque.deepcopy.";
+      func_name += opaque->opaque_type_name;
+      FuncObj *func = Func::GetGlobal(func_name.c_str(), true);
+      if (func == nullptr) {
+        MLC_THROW(ValueError) << "Cannot deepcopy `mlc.Opaque` of type: " << opaque->opaque_type_name
+                              << "; Use `mlc.Func.register(\"" << func_name
+                              << "\")(deepcopy_func)` to register a deepcopy method";
+      }
+      ret = (*func)(object);
     } else {
       fields.clear();
       VisitFields(object, type_info, Copier{&orig2copy, &fields});
@@ -1627,8 +1635,8 @@ inline Any Deserialize(const char *json_str, int64_t json_str_len, FuncObj *fn_o
               MLC_THROW(ValueError) << "Invalid reference when parsing type `" << type_keys[json_type_index]
                                     << "`: referring #" << k << " at #" << i << ". v = " << value;
             }
-          } else if (arg.type_index == kMLCList) {
-            (*list)[j] = invoke_init(arg.operator UList());
+          } else if (UListObj *arg_list = arg.as<UListObj>()) {
+            (*list)[j] = invoke_init(UList(arg_list));
           } else if (arg.type_index == kMLCStr || arg.type_index == kMLCBool || arg.type_index == kMLCFloat ||
                      arg.type_index == kMLCNone) {
             // Do nothing

python/mlc/core/opaque.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+import copy
 from collections.abc import Callable
-from typing import Any
+from typing import Any, Literal
 
 from mlc._cython import Ptr, c_class_core, func_register, opaque_init, register_opauqe_type
 
@@ -18,17 +19,33 @@ def __init__(self, instance: Any) -> None:
     @staticmethod
     def register(
         ty: type,
-        eq_s: Callable | None = None,
-        hash_s: Callable | None = None,
+        eq_s: Callable | Literal["default"] | None = "default",
+        hash_s: Callable | Literal["default"] | None = "default",
+        deepcopy: Callable | Literal["default"] | None = "default",
     ) -> None:
         register_opauqe_type(ty)
         name = ty.__module__ + "." + ty.__name__
-        if eq_s is not None:
-            assert callable(eq_s)
+
+        if isinstance(eq_s, str) and eq_s == "default":
+            func_register(f"Opaque.eq_s.{name}", False, lambda a, b: a == b)
+        elif callable(eq_s):
             func_register(f"Opaque.eq_s.{name}", False, eq_s)
-        if hash_s is not None:
-            assert callable(hash_s)
+        else:
+            assert eq_s is None, "eq_s must be a callable, a literal 'default', or None"
+
+        if isinstance(hash_s, str) and hash_s == "default":
+            func_register(f"Opaque.hash_s.{name}", False, lambda a: hash(a))
+        elif callable(hash_s):
             func_register(f"Opaque.hash_s.{name}", False, hash_s)
+        else:
+            assert hash_s is None, "hash_s must be a callable, a literal 'default', or None"
+
+        if isinstance(deepcopy, str) and deepcopy == "default":
+            func_register(f"Opaque.deepcopy.{name}", False, lambda a: copy.deepcopy(a))
+        elif callable(deepcopy):
+            func_register(f"Opaque.deepcopy.{name}", False, deepcopy)
+        else:
+            assert deepcopy is None, "deepcopy must be a callable, a literal 'default', or None"
 
 
 def _default_serialize(opaques: list[Any]) -> str:

tests/python/test_core_opaque.py

@@ -1,3 +1,4 @@
+import copy
 import json
 from typing import Any
 
@@ -17,19 +18,15 @@ def __init__(self, a: int) -> None:
     def __call__(self, x: int) -> int:
         return x + self.a
 
-    def eq_s(self, other: Any) -> bool:
+    def __eq__(self, other: Any) -> bool:
         return isinstance(self, MyType) and isinstance(other, MyType) and self.a == other.a
 
-    def hash_s(self) -> int:
+    def __hash__(self) -> int:
         assert isinstance(self, MyType)
         return hash((MyType, self.a))
 
 
-mlc.Opaque.register(
-    MyType,
-    eq_s=MyType.eq_s,
-    hash_s=MyType.hash_s,
-)
+mlc.Opaque.register(MyType)
 
 
 @mlc.dataclasses.py_class(structure="bind")
@@ -124,3 +121,13 @@ def test_opaque_serialize_with_alias() -> None:
     assert obj_2.field[3].a == 30
     assert obj_2.field[4].a == 20
     assert obj_2.field[5].a == 10
+
+
+def test_opaque_deepcopy() -> None:
+    a = MyType(a=10)
+    obj_1 = Wrapper(field=a)
+    obj_2 = copy.deepcopy(obj_1)
+    assert isinstance(obj_2.field, MyType)
+    assert obj_2.field.a == 10
+    assert obj_1 is not obj_2
+    assert obj_1.field is not obj_2.field