diff --git a/CMakeLists.txt b/CMakeLists.txt index d0cb8f08d..842bd13bc 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -76,6 +76,7 @@ set(_tvm_ffi_extra_objs_sources "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/structural_equal.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/structural_hash.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/structural_visit.cc" + "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/structural_map.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/visit_error_context.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/json_parser.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/json_writer.cc" diff --git a/include/tvm/ffi/extra/structural_map.h b/include/tvm/ffi/extra/structural_map.h new file mode 100644 index 000000000..da4b35703 --- /dev/null +++ b/include/tvm/ffi/extra/structural_map.h @@ -0,0 +1,613 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +/*! + * \file tvm/ffi/extra/structural_map.h + * \brief Structural mapping and in-place mutation API. + */ +#ifndef TVM_FFI_EXTRA_STRUCTURAL_MAP_H_ +#define TVM_FFI_EXTRA_STRUCTURAL_MAP_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tvm { +namespace ffi { + +class StructuralMapperObj; + +/*! + * \brief ABI of structural transformation hooks and \ref StructuralMapperVTable callbacks. + * + * The callback receives the active mapper as a non-owning pointer and an owning ``Any`` value. It + * returns raw ``TVMFFIAny`` storage containing either the transformed value or an Error. The input + * ownership slot and returned storage are transferred across the ABI boundary by move. + */ +using FStructuralTransform = TVMFFIAny (*)(StructuralMapperObj* mapper, Any value) noexcept; + +namespace details { + +/*! + * \brief Move a structural transformation result to raw ABI storage and annotate failures. + * + * \param result The transformation result to move into raw ABI storage. + * \param error_context The value retained by the current dispatch frame for error reporting. + * \return Raw ``TVMFFIAny`` storing the success value or Error. + */ +TVM_FFI_INLINE static TVMFFIAny MoveStructuralTransformResultToTVMFFIAny( + Expected result, const Any& error_context) noexcept { + if (TVM_FFI_PREDICT_FALSE(result.is_err())) { + if (error_context.type_index() >= TypeIndex::kTVMFFIStaticObjectBegin) { + Error err = result.error(); + UpdateVisitErrorContext(err, error_context.cast()); + } + } + return ExpectedUnsafe::MoveToTVMFFIAny(std::move(result)); +} + +// Dispatch a type-specific structural map or in-place mutation hook. +TVM_FFI_INLINE static Expected DispatchTypeAttrHookExpected( + StructuralMapperObj* mapper, Any value, AnyView attr, std::string_view attr_name) noexcept; + +// Copy and structurally map the reflected fields of an object-backed value. +TVM_FFI_INLINE static Expected MapReflectedFieldsExpected(StructuralMapperObj* mapper, + Any value) noexcept; + +// Structurally transform the reflected fields of an object-backed value in place. +TVM_FFI_INLINE static Expected InplaceMutateReflectedFieldsExpected( + StructuralMapperObj* mapper, Any value) noexcept; + +} // namespace details + +/*! + * \brief VTable ABI for \ref StructuralMapper dispatch. This function table provides a stable ABI + * for the map and in-place mutation methods. + */ +struct StructuralMapperVTable { + /*! + * \brief Select mapping or in-place mutation for a value. + * \param mapper The active structural mapper. + * \param value The owning value to transform by map or in-place mutation. + * \return Raw ``TVMFFIAny`` carrying the transformed value or Error. + */ + FStructuralTransform map_or_inplace_mutate = nullptr; + /*! + * \brief Map a value without intentionally mutating the source. + * \param mapper The active structural mapper. + * \param value The owning value to map. + * \return Raw ``TVMFFIAny`` carrying the mapped value or Error. + */ + FStructuralTransform map = nullptr; + /*! + * \brief Transform a value using the explicit in-place mutation path. + * \param mapper The active structural mapper. + * \param value The owning value to transform in place. + * \return Raw ``TVMFFIAny`` carrying the transformed value or Error. + */ + FStructuralTransform inplace_mutate = nullptr; +}; + +/*! + * \brief Object node of a structural mapper. + * + * A structural mapper recursively transforms values through a manually supplied ABI vtable. + * The default behavior supports type-specific hooks and reflected-field fallback. Values + * are accepted by value so ownership can be transferred without adding persistent references and + * so the default combined operation can select in-place mutation from logical uniqueness. + */ +class StructuralMapperObj : public Object { + public: + /*! \brief Construct the default structural mapper. */ + StructuralMapperObj() : StructuralMapperObj(VTable()) {} + + /*! + * \brief Transform a value, selecting mapping or in-place mutation through the mapper vtable. + * + * \param value The value and one ownership slot to transfer into the transformation. + * \return The transformed value. + * + * \note This method throws the returned Error on failure. Use + * \ref MapOrInplaceMutateExpected for exception-free propagation. + */ + TVM_FFI_INLINE Any MapOrInplaceMutate(Any value) { + return MapOrInplaceMutateExpected(std::move(value)).value(); + } + + /*! + * \brief Exception-free form of \ref MapOrInplaceMutate. + * + * \param value The value and one ownership slot to transfer into the transformation. + * \return The transformed value, or an Error if transformation failed. + */ + TVM_FFI_INLINE Expected MapOrInplaceMutateExpected(Any value) noexcept { + return details::ExpectedUnsafe::MoveFromTVMFFIAny( + (*vtable_->map_or_inplace_mutate)(this, std::move(value))); + } + + /*! + * \brief Apply the default map-or-in-place selection logic. + * + * \param value The value to transform by map or in-place mutation. + * \return The transformed value. + */ + TVM_FFI_INLINE Any DefaultMapOrInplaceMutate(Any value) { + return DefaultMapOrInplaceMutateExpected(std::move(value)).value(); + } + + /*! + * \brief Exception-free form of \ref DefaultMapOrInplaceMutate. + * + * \param value The value to transform by map or in-place mutation. + * \return The transformed value, or an Error if selection or transformation failed. + */ + TVM_FFI_INLINE Expected DefaultMapOrInplaceMutateExpected(Any value) noexcept { + return DefaultMapOrInplaceMutateExpected(std::move(value), false); + } + + /*! + * \brief Map a value through the mapper vtable. + * + * \param value The value to map. + * \return The mapped value. + * + * \note This method throws the returned Error on failure. Use \ref MapExpected for + * exception-free propagation. + */ + TVM_FFI_INLINE Any Map(Any value) { return MapExpected(std::move(value)).value(); } + + /*! + * \brief Exception-free form of \ref Map. + * + * \param value The value to map. + * \return The mapped value, or an Error if mapping failed. + */ + TVM_FFI_INLINE Expected MapExpected(Any value) noexcept { + return details::ExpectedUnsafe::MoveFromTVMFFIAny((*vtable_->map)(this, std::move(value))); + } + + /*! + * \brief Apply the default structural map with copy-on-write behavior. + * + * \param value The value to map. + * \return The mapped value. + */ + TVM_FFI_INLINE Any DefaultMap(Any value) { return DefaultMapExpected(std::move(value)).value(); } + + /*! + * \brief Exception-free form of \ref DefaultMap. + * + * \param value The value to map. + * \return The mapped value, or an Error if hook dispatch, copying, or field mapping failed. + */ + TVM_FFI_INLINE Expected DefaultMapExpected(Any value) noexcept { + int32_t type_index = value.type_index(); + static reflection::TypeAttrColumn column(reflection::type_attr::kStructuralMap); + AnyView attr = column[type_index]; + if (attr.type_index() != TypeIndex::kTVMFFINone) { + return details::DispatchTypeAttrHookExpected(this, std::move(value), attr, + reflection::type_attr::kStructuralMap); + } + if (type_index < TypeIndex::kTVMFFIStaticObjectBegin) { + return value; + } + return details::MapReflectedFieldsExpected(this, std::move(value)); + } + + /*! + * \brief Transform a value through the explicit in-place mutation vtable entry. + * + * \param value The value to transform in place. + * \return The transformed value. + */ + TVM_FFI_INLINE Any InplaceMutate(Any value) { + return InplaceMutateExpected(std::move(value)).value(); + } + + /*! + * \brief Exception-free form of \ref InplaceMutate. + * + * \param value The value to transform in place. + * \return The transformed value, or an Error if transformation failed. + */ + TVM_FFI_INLINE Expected InplaceMutateExpected(Any value) noexcept { + return details::ExpectedUnsafe::MoveFromTVMFFIAny( + (*vtable_->inplace_mutate)(this, std::move(value))); + } + + /*! + * \brief Apply the default structural in-place mutation. + * + * \param value The value to transform in place. + * \return The transformed value. + */ + TVM_FFI_INLINE Any DefaultInplaceMutate(Any value) { + return DefaultInplaceMutateExpected(std::move(value)).value(); + } + + /*! + * \brief Exception-free form of \ref DefaultInplaceMutate. + * + * \param value The value to transform in place. + * \return The transformed value, or an Error if hook dispatch or reflected mutation failed. + */ + TVM_FFI_INLINE Expected DefaultInplaceMutateExpected(Any value) noexcept { + int32_t type_index = value.type_index(); + static reflection::TypeAttrColumn column(reflection::type_attr::kStructuralInplaceMutate); + AnyView attr = column[type_index]; + if (attr.type_index() != TypeIndex::kTVMFFINone) { + return details::DispatchTypeAttrHookExpected(this, std::move(value), attr, + reflection::type_attr::kStructuralInplaceMutate); + } + if (type_index < TypeIndex::kTVMFFIStaticObjectBegin) { + return value; + } + return details::InplaceMutateReflectedFieldsExpected(this, std::move(value)); + } + + /// \cond Doxygen_Suppress + static constexpr const bool _type_mutable = true; + TVM_FFI_DECLARE_OBJECT_INFO("ffi.StructuralMapper", StructuralMapperObj, Object); + /// \endcond + + protected: + /*! + * \brief Construct a structural mapper subclass with a custom dispatch vtable. + * + * \param vtable The non-null dispatch table for this mapper. + * + * \note This constructor is for internal subclasses. The vtable and its + * ``map_or_inplace_mutate`` callback must be valid for the lifetime of the mapper. + */ + explicit StructuralMapperObj(const StructuralMapperVTable* vtable) : vtable_(vtable) {} + + /*! + * \brief Apply default combined transformation with a customized uniqueness decision. + * + * \param value The value to transform. + * \param can_inplace_mutate Whether the caller has already established logical uniqueness before + * adding temporary internal owners. + * \return The transformed value. + * + * \note This protected overload is for ownership-aware internal recursion. Passing ``true`` is + * invalid when the object is genuinely shared. + */ + TVM_FFI_INLINE Any DefaultMapOrInplaceMutate(Any value, bool can_inplace_mutate) { + return DefaultMapOrInplaceMutateExpected(std::move(value), can_inplace_mutate).value(); + } + + /*! + * \brief Exception-free form of the default transformation with a customized uniqueness decision. + * + * \param value The value to transform. + * \param can_inplace_mutate Whether the caller has already established logical uniqueness before + * adding temporary internal owners. + * \return The transformed value, or an Error if validation or transformation failed. + */ + TVM_FFI_INLINE Expected DefaultMapOrInplaceMutateExpected(Any value, + bool can_inplace_mutate) noexcept { + if (const Object* obj = value.as()) { + // check both map and inplace_mutate hooks are defined + static reflection::TypeAttrColumn map_column(reflection::type_attr::kStructuralMap); + static reflection::TypeAttrColumn inplace_mutate_column( + reflection::type_attr::kStructuralInplaceMutate); + int32_t type_index = value.type_index(); + AnyView map_attr = map_column[type_index]; + AnyView inplace_mutate_attr = inplace_mutate_column[type_index]; + bool has_map = map_attr.type_index() != TypeIndex::kTVMFFINone; + bool has_inplace_mutate = inplace_mutate_attr.type_index() != TypeIndex::kTVMFFINone; + if (has_map != has_inplace_mutate) { + return Unexpected(Error("TypeError", + "One of " + std::string(reflection::type_attr::kStructuralMap) + + " and " + + std::string(reflection::type_attr::kStructuralInplaceMutate) + + " is undefined, should provide both of them.", + "")); + } + if (obj->unique() || can_inplace_mutate) { + return InplaceMutateExpected(std::move(value)); + } + } + return MapExpected(std::move(value)); + } + + // Grant reflected in-place recursion access to the ownership-aware protected overload. + friend Expected details::InplaceMutateReflectedFieldsExpected(StructuralMapperObj* mapper, + Any value) noexcept; + + /*! + * \brief Required ABI dispatch table. \ref StructuralMapperVTable + * It must never be null on a constructed mapper. + */ + const StructuralMapperVTable* vtable_ = nullptr; + + private: + /*! + * \brief Return the vtable used by the default structural mapper. + * \return Pointer to the static default vtable. + */ + static const StructuralMapperVTable* VTable() { + static const StructuralMapperVTable vtable{ + &StructuralMapperObj::DispatchMapOrInplaceMutate, + &StructuralMapperObj::DispatchMap, + &StructuralMapperObj::DispatchInplaceMutate, + }; + return &vtable; + } + + /*! + * \brief Dispatch the default transformation from the ABI vtable. + * + * \param mapper The active structural mapper. + * \param value The owning value to transform. + * \return Raw ``TVMFFIAny`` storing the transformed value or Error. + * + * \note Validation errors raised before a branch is selected are returned without adding a + * visit-context frame at this forwarding layer to avoid duplicated error context. + */ + static TVMFFIAny DispatchMapOrInplaceMutate(StructuralMapperObj* mapper, Any value) noexcept { + auto result = mapper->DefaultMapOrInplaceMutateExpected(std::move(value)); + return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result)); + } + + /*! + * \brief Dispatch default mapping from the ABI vtable. + * + * \param mapper The active structural mapper. + * \param value The owning value to map. + * \return Raw ``TVMFFIAny`` storing the mapped value or Error. + */ + static TVMFFIAny DispatchMap(StructuralMapperObj* mapper, Any value) noexcept { + Any error_context = value; + auto result = mapper->DefaultMapExpected(std::move(value)); + return details::MoveStructuralTransformResultToTVMFFIAny(std::move(result), error_context); + } + + /*! + * \brief Dispatch default in-place mutation from the ABI vtable. + * + * \param mapper The active structural mapper. + * \param value The owning value to transform in place. + * \return Raw ``TVMFFIAny`` storing the transformed value or Error. + */ + static TVMFFIAny DispatchInplaceMutate(StructuralMapperObj* mapper, Any value) noexcept { + Any error_context = value; + auto result = mapper->DefaultInplaceMutateExpected(std::move(value)); + return details::MoveStructuralTransformResultToTVMFFIAny(std::move(result), error_context); + } +}; + +/*! + * \brief ObjectRef wrapper of \ref StructuralMapperObj. + * + * \sa StructuralMapperObj + */ +class StructuralMapper : public ObjectRef { + public: + /*! \brief Construct the default structural mapper. */ + StructuralMapper() : ObjectRef(make_object()) {} + + /*! + * \brief Construct from an existing mapper object pointer. + * \param n The object pointer to wrap. + */ + explicit StructuralMapper(ObjectPtr n) : ObjectRef(std::move(n)) {} + + /// \cond Doxygen_Suppress + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StructuralMapper, ObjectRef, StructuralMapperObj); + /// \endcond +}; + +namespace details { + +/*! + * \brief Dispatch a type-specific structural transformation hook. + * + * \param mapper The active structural mapper. + * \param value The value to pass to the hook. + * \param attr The registered type attribute value. + * \param attr_name The attribute name used in type errors. + * \return The hook result, or an Error for hook failure or an invalid attribute value. + */ +TVM_FFI_INLINE static Expected DispatchTypeAttrHookExpected( + StructuralMapperObj* mapper, Any value, AnyView attr, std::string_view attr_name) noexcept { + // case 1: Type-specific override registered as an opaque ABI function pointer. + if (attr.type_index() == TypeIndex::kTVMFFIOpaquePtr) { + auto* hook = reinterpret_cast(attr.cast()); + return details::ExpectedUnsafe::MoveFromTVMFFIAny((*hook)(mapper, std::move(value))); + } + + // case 2: Type-specific override registered as an ffi::Function. + if (attr.type_index() == TypeIndex::kTVMFFIFunction) { + return attr.cast().CallExpected(mapper, value); + } + + return Unexpected( + Error("TypeError", + std::string(attr_name) + " must be an opaque function pointer or ffi.Function", "")); +} + +/*! + * \brief Transform every reflected structural field of an object. + * + * \tparam Callback A callable compatible with ``Expected(const Any&)``. + * \param value The original object-backed value. + * \param result The shallow-copy result in copy-on-write mode; ignored in in-place mode. + * \param copy_on_write Whether to transform a distinct shallow copy instead of \p value. + * \param callback The recursive field transformation callback. + * \return The original value, transformed copy, or in-place result; otherwise an Error. + */ +template +TVM_FFI_INLINE static Expected TransformReflectedFieldsExpected(Any value, + Expected result, + bool copy_on_write, + Callback callback) noexcept { + const Object* obj = value.as(); + if (!copy_on_write) { + // In-place mode transfers the input ownership slot into the result container. + result = std::exchange(value, Any()); + } + const Any& result_value = details::ExpectedUnsafe::GetData(result); + Object* new_obj = const_cast(result_value.as()); + if (copy_on_write) { + // Mapping requires a distinct target so partially applied updates cannot mutate the source. + if (TVM_FFI_PREDICT_FALSE(new_obj == nullptr || result.type_index() != value.type_index() || + new_obj == obj)) { + return Unexpected( + Error("TypeError", + "Shallow copy callback must return a distinct object with the same type as its " + "input", + "")); + } + } + const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(new_obj->type_index()); + + bool field_changed = false; + reflection::ForEachFieldInfoWithEarlyStop( + type_info, [&](const TVMFFIFieldInfo* field_info) -> bool { + if (field_info->flags & kTVMFFIFieldFlagBitMaskSEqHashIgnore) { + return false; + } + + Any field_value; + const void* field_addr = reinterpret_cast(new_obj) + field_info->offset; + int ret_code = field_info->getter(const_cast(field_addr), + reinterpret_cast(&field_value)); + if (TVM_FFI_PREDICT_FALSE(ret_code != 0)) { + result = Unexpected(details::MoveFromSafeCallRaised()); + return true; + } + + // TODO(kathrync): add WithDefRegionKind + Expected transformed_field = callback(field_value); + if (TVM_FFI_PREDICT_FALSE(transformed_field.is_err())) { + result = std::move(transformed_field); + return true; + } + const Any& new_field = details::ExpectedUnsafe::GetData(transformed_field); + if (field_value.same_as(new_field)) { + return false; + } + + if (TVM_FFI_PREDICT_FALSE(field_info->setter == nullptr)) { + result = Unexpected( + Error("TypeError", + "Cannot structurally map field `" + + std::string(field_info->name.data, field_info->name.size) + "` of type `" + + std::string(type_info->type_key.data, type_info->type_key.size) + + "` because it does not define a setter", + "")); + return true; + } + + void* new_field_addr = reinterpret_cast(new_obj) + field_info->offset; + ret_code = reflection::CallFieldSetter(field_info, new_field_addr, + reinterpret_cast(&new_field)); + if (TVM_FFI_PREDICT_FALSE(ret_code != 0)) { + result = Unexpected(details::MoveFromSafeCallRaised()); + return true; + } + field_changed = true; + return false; + }); + + if (TVM_FFI_PREDICT_FALSE(result.is_err())) { + return result; + } + if (copy_on_write && !field_changed) { + return value; + } + return result; +} + +/*! + * \brief Map the reflected structural fields of an object-backed value. + * + * \param mapper The active structural mapper. + * \param value The object-backed value to map. + * \return The original value when no field changes, a transformed shallow copy otherwise, or an + * Error if copying or mapping failed. + */ +TVM_FFI_INLINE static Expected MapReflectedFieldsExpected(StructuralMapperObj* mapper, + Any value) noexcept { + const Object* obj = value.as(); + int32_t type_index = obj->type_index(); + + static reflection::TypeAttrColumn column(reflection::type_attr::kShallowCopy); + AnyView attr = column[type_index]; + if (TVM_FFI_PREDICT_FALSE(attr.type_index() != TypeIndex::kTVMFFIFunction)) { + return Unexpected( + Error("TypeError", + std::string(reflection::type_attr::kShallowCopy) + " must be an ffi.Function", "")); + } + + Expected result = attr.cast().CallExpected(value); + if (TVM_FFI_PREDICT_FALSE(result.is_err())) { + return result; + } + + return TransformReflectedFieldsExpected( + std::move(value), std::move(result), true, + [mapper](const Any& field_value) noexcept -> Expected { + return mapper->MapExpected(field_value); + }); +} + +/*! + * \brief Transform the reflected structural fields of an object-backed value in place. + * + * \param mapper The active structural mapper. + * \param value The object-backed value whose fields should be transformed in place. + * \return The input object after field transformation, or an Error. Mutations made before an Error + * are not rolled back. + */ +TVM_FFI_INLINE static Expected InplaceMutateReflectedFieldsExpected( + StructuralMapperObj* mapper, Any value) noexcept { + return TransformReflectedFieldsExpected( + std::move(value), Any(), false, [mapper](const Any& field_value) noexcept -> Expected { + const Object* field_obj = field_value.as(); + // A logically unique object field has two references here: one from its parent and one + // from field_value, which was created by the reflection getter. + bool can_inplace_mutate = field_obj != nullptr && field_obj->use_count() == 2; + return mapper->DefaultMapOrInplaceMutateExpected(field_value, can_inplace_mutate); + }); +} + +} // namespace details +} // namespace ffi +} // namespace tvm + +#endif // TVM_FFI_EXTRA_STRUCTURAL_MAP_H_ diff --git a/include/tvm/ffi/reflection/accessor.h b/include/tvm/ffi/reflection/accessor.h index d0aa734cc..889574122 100644 --- a/include/tvm/ffi/reflection/accessor.h +++ b/include/tvm/ffi/reflection/accessor.h @@ -495,6 +495,49 @@ inline constexpr const char* kSEqual = "__s_equal__"; * visit structural children. */ inline constexpr const char* kStructuralVisit = "__s_visit__"; + +/*! + * \brief Custom structural map hook used by ``StructuralMapper``. + * + * The hook receives the active mapper and the input value. It should recursively map structural + * children through the mapper, return the input unchanged if no fields are changed, and avoid + * intentionally mutating the source object. An opaque hook owns its ``Any`` argument; an + * ``ffi::Function`` receives a borrowed packed view kept alive for the duration of the call. + * + * Value type: either an opaque function pointer to a C++ structural map hook + * + * ``TVMFFIAny (*)(StructuralMapperObj* mapper, Any value) noexcept`` + * + * returning raw ``Expected`` storage, or an ``ffi::Function`` with signature + * + * ``(StructuralMapper mapper, Any value) -> Any``. + * + * A type used with the default combined map-or-in-place operation must define this attribute and + * ``kStructuralInplaceMutate`` together, or leave both undefined. + */ +inline constexpr const char* kStructuralMap = "__s_map__"; +/*! + * \brief Custom structural in-place mutation hook used by ``StructuralMapper``. + * + * The hook receives the active mapper and the input value. It may mutate the source object and + * should normally return that same object. An opaque hook owns its ``Any`` argument; an + * ``ffi::Function`` receives a borrowed packed view kept alive for the duration of the call. The + * explicit in-place path does not verify uniqueness and does not roll back changes when a later + * transformation fails. + * + * Value type: either an opaque function pointer to a C++ structural in-place mutation hook + * + * ``TVMFFIAny (*)(StructuralMapperObj* mapper, Any value) noexcept`` + * + * returning raw ``Expected`` storage, or an ``ffi::Function`` with signature + * + * ``(StructuralMapper mapper, Any value) -> Any``. + * + * A type used with the default combined map-or-in-place operation must define this attribute and + * ``kStructuralMap`` together, or leave both undefined. + */ +inline constexpr const char* kStructuralInplaceMutate = "__s_inplace_mutate__"; + /*! * \brief Serialize object data to a JSON-compatible ``Map``. * diff --git a/src/ffi/extra/structural_map.cc b/src/ffi/extra/structural_map.cc new file mode 100644 index 000000000..dd3c3bc5a --- /dev/null +++ b/src/ffi/extra/structural_map.cc @@ -0,0 +1,49 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +/*! + * \file src/ffi/extra/structural_map.cc + * \brief Structural map registration. + */ +#include +#include +#include + +namespace tvm { +namespace ffi { + +// --------------------------------------------------------------------------- +// Static registration. +// --------------------------------------------------------------------------- + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef().def( + refl::init<>(), "Constructor that creates a default structural mapper"); + refl::GlobalDef() + .def("ffi.StructuralMapper", []() { return StructuralMapper(); }) + .def_method("ffi.StructuralMapperMapOrInplaceMutate", + &StructuralMapperObj::MapOrInplaceMutate) + .def_method("ffi.StructuralMapperMap", &StructuralMapperObj::Map) + .def_method("ffi.StructuralMapperInplaceMutate", &StructuralMapperObj::InplaceMutate); + refl::EnsureTypeAttrColumn(refl::type_attr::kStructuralMap); + refl::EnsureTypeAttrColumn(refl::type_attr::kStructuralInplaceMutate); +} + +} // namespace ffi +} // namespace tvm diff --git a/src/ffi/testing/testing.cc b/src/ffi/testing/testing.cc index 4447900a1..1c06572ff 100644 --- a/src/ffi/testing/testing.cc +++ b/src/ffi/testing/testing.cc @@ -37,6 +37,7 @@ #include #include +#include #include #include @@ -379,6 +380,162 @@ class TestEqWithoutHash : public ObjectRef { TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TestEqWithoutHash, ObjectRef, TestEqWithoutHashObj); }; +enum class StructuralMapTestAction : int32_t { + kKeep, + kChange, + kError, +}; + +class StructuralMapperObj; +using FStructuralMapTestTransform = TVMFFIAny (*)(StructuralMapperObj*, Any) noexcept; + +Array& StructuralMapTestTrace() { + static Array trace; + return trace; +} + +class StructuralMapTestLeafObj : public Object { + public: + int64_t value; + StructuralMapTestAction action; + + StructuralMapTestLeafObj(int64_t value, StructuralMapTestAction action) + : value(value), action(action) {} + + static TVMFFIAny StructuralMap(StructuralMapperObj* mapper, Any value) noexcept; + static TVMFFIAny StructuralInplaceMutate(StructuralMapperObj* mapper, Any value) noexcept; + + static constexpr bool _type_mutable = true; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("testing.StructuralMapLeaf", StructuralMapTestLeafObj, Object); +}; + +class StructuralMapTestLeaf : public ObjectRef { + public: + StructuralMapTestLeaf(int64_t value, StructuralMapTestAction action) { + data_ = make_object(value, action); + } + + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(StructuralMapTestLeaf, ObjectRef, + StructuralMapTestLeafObj); +}; + +class StructuralMapTestFunctionLeafObj : public Object { + public: + int64_t value; + StructuralMapTestAction action; + + StructuralMapTestFunctionLeafObj(int64_t value, StructuralMapTestAction action) + : value(value), action(action) {} + + static Any StructuralMap(ObjectRef, Any value); + static Any StructuralInplaceMutate(ObjectRef, Any value); + + static constexpr bool _type_mutable = true; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("testing.StructuralMapFunctionLeaf", + StructuralMapTestFunctionLeafObj, Object); +}; + +class StructuralMapTestFunctionLeaf : public ObjectRef { + public: + StructuralMapTestFunctionLeaf(int64_t value, StructuralMapTestAction action) { + data_ = make_object(value, action); + } + + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(StructuralMapTestFunctionLeaf, ObjectRef, + StructuralMapTestFunctionLeafObj); +}; + +// NOLINTNEXTLINE(bugprone-exception-escape) +TVMFFIAny StructuralMapTestLeafObj::StructuralMap(StructuralMapperObj*, Any value) noexcept { + const StructuralMapTestLeafObj* leaf = value.as(); + StructuralMapTestTrace().push_back("map:" + std::to_string(leaf->value)); + + if (leaf->action == StructuralMapTestAction::kError) { + Expected result = Unexpected(Error("ValueError", "structural map leaf failed", "")); + return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result)); + } + if (leaf->action == StructuralMapTestAction::kChange) { + Expected result = + Any(StructuralMapTestLeaf(leaf->value + 1, StructuralMapTestAction::kKeep)); + return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result)); + } + + Expected result = std::move(value); + return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result)); +} + +// NOLINTNEXTLINE(bugprone-exception-escape) +TVMFFIAny StructuralMapTestLeafObj::StructuralInplaceMutate(StructuralMapperObj*, + Any value) noexcept { + StructuralMapTestLeafObj* leaf = value.cast(); + StructuralMapTestTrace().push_back("inplace:" + std::to_string(leaf->value)); + + if (leaf->action == StructuralMapTestAction::kError) { + Expected result = Unexpected(Error("ValueError", "structural map leaf failed", "")); + return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result)); + } + if (leaf->action == StructuralMapTestAction::kChange) { + ++leaf->value; + leaf->action = StructuralMapTestAction::kKeep; + } + + Expected result = std::move(value); + return details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result)); +} + +// NOLINTNEXTLINE(performance-unnecessary-value-param) +Any StructuralMapTestFunctionLeafObj::StructuralMap(ObjectRef, Any value) { + const StructuralMapTestFunctionLeafObj* leaf = value.as(); + StructuralMapTestTrace().push_back("function-map:" + std::to_string(leaf->value)); + + if (leaf->action == StructuralMapTestAction::kError) { + TVM_FFI_THROW(ValueError) << "structural map function leaf failed"; + } + if (leaf->action == StructuralMapTestAction::kChange) { + return StructuralMapTestFunctionLeaf(leaf->value + 1, StructuralMapTestAction::kKeep); + } + return value; +} + +// NOLINTNEXTLINE(performance-unnecessary-value-param) +Any StructuralMapTestFunctionLeafObj::StructuralInplaceMutate(ObjectRef, Any value) { + StructuralMapTestFunctionLeafObj* leaf = value.cast(); + StructuralMapTestTrace().push_back("function-inplace:" + std::to_string(leaf->value)); + + if (leaf->action == StructuralMapTestAction::kError) { + TVM_FFI_THROW(ValueError) << "structural map function leaf failed"; + } + if (leaf->action == StructuralMapTestAction::kChange) { + ++leaf->value; + leaf->action = StructuralMapTestAction::kKeep; + } + return value; +} + +ObjectRef MakeStructuralMapTestLeaf(int64_t value, int32_t action) { + return StructuralMapTestLeaf(value, static_cast(action)); +} + +ObjectRef MakeStructuralMapTestFunctionLeaf(int64_t value, int32_t action) { + return StructuralMapTestFunctionLeaf(value, static_cast(action)); +} + +int64_t GetStructuralMapTestLeafValue(const ObjectRef& value) { + if (const auto* leaf = value.as()) { + return leaf->value; + } + if (const auto* leaf = value.as()) { + return leaf->value; + } + TVM_FFI_THROW(TypeError) << "Expected a structural-map test leaf"; +} + +Array GetStructuralMapTestTrace() { return StructuralMapTestTrace(); } + +void ClearStructuralMapTestTrace() { StructuralMapTestTrace().clear(); } + // NOLINTNEXTLINE(performance-unnecessary-value-param) TVM_FFI_NO_INLINE void TestRaiseError(String kind, String msg) { // keep name and no liner for testing backtrace @@ -535,6 +692,25 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def_ro("key", &TestEqWithoutHashObj::key) .def_ro("label", &TestEqWithoutHashObj::label); + refl::EnsureTypeAttrColumn(refl::type_attr::kStructuralMap); + refl::EnsureTypeAttrColumn(refl::type_attr::kStructuralInplaceMutate); + + refl::ObjectDef().def_rw("value", &StructuralMapTestLeafObj::value); + refl::TypeAttrDef() + .attr(refl::type_attr::kStructuralMap, + reinterpret_cast( + static_cast(&StructuralMapTestLeafObj::StructuralMap))) + .attr(refl::type_attr::kStructuralInplaceMutate, + reinterpret_cast(static_cast( + &StructuralMapTestLeafObj::StructuralInplaceMutate))); + + refl::ObjectDef().def_rw( + "value", &StructuralMapTestFunctionLeafObj::value); + refl::TypeAttrDef() + .def(refl::type_attr::kStructuralMap, &StructuralMapTestFunctionLeafObj::StructuralMap) + .def(refl::type_attr::kStructuralInplaceMutate, + &StructuralMapTestFunctionLeafObj::StructuralInplaceMutate); + refl::GlobalDef() .def("testing.test_raise_error", TestRaiseError) .def("testing.add_one", [](int x) { return x + 1; }) @@ -616,6 +792,13 @@ TVM_FFI_STATIC_INIT_BLOCK() { return result; }); // NOLINTEND(performance-unnecessary-value-param) + + refl::GlobalDef() + .def("testing.make_structural_map_leaf", MakeStructuralMapTestLeaf) + .def("testing.make_structural_map_function_leaf", MakeStructuralMapTestFunctionLeaf) + .def("testing.structural_map_leaf_value", GetStructuralMapTestLeafValue) + .def("testing.structural_mapper_trace", GetStructuralMapTestTrace) + .def("testing.structural_mapper_clear_trace", ClearStructuralMapTestTrace); } } // namespace ffi diff --git a/tests/cpp/extra/test_structural_map.cc b/tests/cpp/extra/test_structural_map.cc new file mode 100644 index 000000000..9e66259e5 --- /dev/null +++ b/tests/cpp/extra/test_structural_map.cc @@ -0,0 +1,428 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "../testing_object.h" + +namespace { + +using namespace tvm::ffi; +using namespace tvm::ffi::testing; + +enum class TransformAction : int32_t { + kKeep, + kChange, + kError, +}; + +using TestMapper = StructuralMapper; + +ObjectRef MakeMapLeaf(int64_t value, TransformAction action = TransformAction::kKeep) { + static Function make_leaf = Function::GetGlobalRequired("testing.make_structural_map_leaf"); + return make_leaf(value, static_cast(action)).cast(); +} + +ObjectRef MakeFunctionLeaf(int64_t value, TransformAction action = TransformAction::kKeep) { + static Function make_leaf = + Function::GetGlobalRequired("testing.make_structural_map_function_leaf"); + return make_leaf(value, static_cast(action)).cast(); +} + +int64_t MapLeafValue(const ObjectRef& value) { + static Function get_value = Function::GetGlobalRequired("testing.structural_map_leaf_value"); + return get_value(value).cast(); +} + +Array MapperTrace() { + static Function get_trace = Function::GetGlobalRequired("testing.structural_mapper_trace"); + return get_trace().cast>(); +} + +void ClearMapperTrace() { + static Function clear_trace = + Function::GetGlobalRequired("testing.structural_mapper_clear_trace"); + clear_trace(); +} + +TestMapper MakeTestMapper() { + ClearMapperTrace(); + return StructuralMapper(); +} + +void ExpectTrace(std::initializer_list expected) { + Array trace = MapperTrace(); + ASSERT_EQ(trace.size(), expected.size()); + size_t i = 0; + for (const char* item : expected) { + EXPECT_EQ(trace[i], item); + ++i; + } +} + +// --------------------------------------------------------------------------- +// StructuralMapper behavior. +// --------------------------------------------------------------------------- + +TEST(StructuralMapper, ReturnsPODValuesUnchanged) { + StructuralMapper mapper; + + Expected mapped = mapper->MapExpected(int64_t{42}); + ASSERT_TRUE(mapped.is_ok()); + EXPECT_EQ(std::move(mapped).value().cast(), 42); + + Expected mutated = mapper->InplaceMutateExpected(3.5); + ASSERT_TRUE(mutated.is_ok()); + EXPECT_EQ(std::move(mutated).value().cast(), 3.5); + + Expected combined = mapper->MapOrInplaceMutateExpected(true); + ASSERT_TRUE(combined.is_ok()); + EXPECT_TRUE(std::move(combined).value().cast()); +} + +TEST(StructuralMapper, MapReturnsOriginalWhenFieldsUnchanged) { + ObjectRef lhs = MakeMapLeaf(1); + TVar rhs("rhs"); + const Object* lhs_addr = lhs.get(); + const Object* rhs_addr = rhs.get(); + TPair root(std::move(lhs), std::move(rhs)); + const Object* root_addr = root.get(); + Any input = std::move(root); + TestMapper mapper = MakeTestMapper(); + + Expected result = mapper->MapExpected(std::move(input)); + + ASSERT_TRUE(result.is_ok()); + TPair mapped = std::move(result).value().cast(); + EXPECT_EQ(mapped.get(), root_addr); + EXPECT_EQ(mapped->lhs.get(), lhs_addr); + EXPECT_EQ(mapped->rhs.get(), rhs_addr); + EXPECT_EQ(mapped.use_count(), 1); + ExpectTrace({"map:1"}); +} + +TEST(StructuralMapper, MapCreatesShallowCopyWhenFieldChanges) { + ObjectRef lhs = MakeMapLeaf(1, TransformAction::kChange); + TVar rhs("rhs"); + TPair root(std::move(lhs), std::move(rhs)); + const Object* root_addr = root.get(); + const Object* lhs_addr = root->lhs.get(); + const Object* rhs_addr = root->rhs.get(); + TestMapper mapper = MakeTestMapper(); + + Expected result = mapper->MapExpected(root); + + ASSERT_TRUE(result.is_ok()); + TPair mapped = std::move(result).value().cast(); + EXPECT_NE(mapped.get(), root_addr); + EXPECT_EQ(root->lhs.get(), lhs_addr); + EXPECT_EQ(MapLeafValue(root->lhs), 1); + EXPECT_NE(mapped->lhs.get(), lhs_addr); + EXPECT_EQ(MapLeafValue(mapped->lhs), 2); + EXPECT_EQ(mapped->rhs.get(), rhs_addr); + EXPECT_TRUE(mapped->rhs.same_as(root->rhs)); + + EXPECT_EQ(root.use_count(), 1); + EXPECT_EQ(mapped.use_count(), 1); + EXPECT_EQ(root->lhs.use_count(), 1); + EXPECT_EQ(mapped->lhs.use_count(), 1); + EXPECT_EQ(root->rhs.use_count(), 2); + ExpectTrace({"map:1"}); +} + +TEST(StructuralMapper, MapCopiesOnlyChangedPath) { + ObjectRef leaf = MakeMapLeaf(1, TransformAction::kChange); + TVar middle_rhs("middle-rhs"); + TPair middle(std::move(leaf), std::move(middle_rhs)); + const TPairObj* middle_addr = middle.get(); + const Object* leaf_addr = middle->lhs.get(); + const Object* middle_rhs_addr = middle->rhs.get(); + + TVar outer_rhs("outer-rhs"); + TPair root(std::move(middle), std::move(outer_rhs)); + const Object* root_addr = root.get(); + const Object* outer_rhs_addr = root->rhs.get(); + TestMapper mapper = MakeTestMapper(); + + Expected result = mapper->MapExpected(root); + + ASSERT_TRUE(result.is_ok()); + TPair mapped = std::move(result).value().cast(); + const TPairObj* mapped_middle = mapped->lhs.as(); + ASSERT_NE(mapped_middle, nullptr); + + EXPECT_NE(mapped.get(), root_addr); + EXPECT_NE(mapped_middle, middle_addr); + EXPECT_NE(mapped_middle->lhs.get(), leaf_addr); + EXPECT_EQ(MapLeafValue(mapped_middle->lhs), 2); + + EXPECT_EQ(middle_addr->lhs.get(), leaf_addr); + EXPECT_EQ(MapLeafValue(middle_addr->lhs), 1); + EXPECT_EQ(mapped_middle->rhs.get(), middle_rhs_addr); + EXPECT_EQ(mapped->rhs.get(), outer_rhs_addr); + ExpectTrace({"map:1"}); +} + +TEST(StructuralMapper, MapOrInplaceMutateUsesInplaceForUniqueRoot) { + ObjectRef lhs = MakeMapLeaf(1, TransformAction::kChange); + ObjectRef rhs = MakeMapLeaf(2); + TPair root(std::move(lhs), std::move(rhs)); + const Object* root_addr = root.get(); + const Object* lhs_addr = root->lhs.get(); + const Object* rhs_addr = root->rhs.get(); + Any input = std::move(root); + TestMapper mapper = MakeTestMapper(); + + Expected result = mapper->MapOrInplaceMutateExpected(std::move(input)); + + ASSERT_TRUE(result.is_ok()); + TPair mapped = std::move(result).value().cast(); + EXPECT_EQ(mapped.get(), root_addr); + EXPECT_EQ(mapped->lhs.get(), lhs_addr); + EXPECT_EQ(MapLeafValue(mapped->lhs), 2); + EXPECT_EQ(mapped->rhs.get(), rhs_addr); + EXPECT_EQ(mapped.use_count(), 1); + EXPECT_EQ(mapped->lhs.use_count(), 1); + EXPECT_EQ(mapped->rhs.use_count(), 1); + ExpectTrace({"inplace:1", "inplace:2"}); +} + +TEST(StructuralMapper, MapOrInplaceMutateUsesMapForSharedRoot) { + ObjectRef lhs = MakeMapLeaf(1, TransformAction::kChange); + ObjectRef rhs = MakeMapLeaf(2); + TPair root(std::move(lhs), std::move(rhs)); + const Object* root_addr = root.get(); + const Object* lhs_addr = root->lhs.get(); + TestMapper mapper = MakeTestMapper(); + + Expected result = mapper->MapOrInplaceMutateExpected(root); + + ASSERT_TRUE(result.is_ok()); + TPair mapped = std::move(result).value().cast(); + EXPECT_NE(mapped.get(), root_addr); + EXPECT_EQ(root->lhs.get(), lhs_addr); + EXPECT_EQ(MapLeafValue(root->lhs), 1); + EXPECT_NE(mapped->lhs.get(), lhs_addr); + EXPECT_EQ(MapLeafValue(mapped->lhs), 2); + EXPECT_TRUE(mapped->rhs.same_as(root->rhs)); + EXPECT_EQ(root.use_count(), 1); + EXPECT_EQ(mapped.use_count(), 1); + ExpectTrace({"map:1", "map:2"}); +} + +TEST(StructuralMapper, ExplicitInplaceMutateUpdatesSharedRoot) { + ObjectRef lhs = MakeMapLeaf(1, TransformAction::kChange); + ObjectRef rhs = MakeMapLeaf(2); + TPair root(std::move(lhs), std::move(rhs)); + TPair alias = root; + const Object* root_addr = root.get(); + const Object* lhs_addr = root->lhs.get(); + TestMapper mapper = MakeTestMapper(); + + Expected result = mapper->InplaceMutateExpected(std::move(root)); + + ASSERT_TRUE(result.is_ok()); + TPair mapped = std::move(result).value().cast(); + EXPECT_EQ(mapped.get(), root_addr); + EXPECT_EQ(alias.get(), root_addr); + EXPECT_EQ(alias->lhs.get(), lhs_addr); + EXPECT_EQ(MapLeafValue(alias->lhs), 2); + EXPECT_EQ(mapped.use_count(), 2); + ExpectTrace({"inplace:1", "inplace:2"}); +} + +void CheckInplaceFieldSelection(bool change_lhs, bool change_rhs) { + SCOPED_TRACE(std::string("change_lhs=") + (change_lhs ? "true" : "false") + + ", change_rhs=" + (change_rhs ? "true" : "false")); + + // Before the getter runs, lhs is shared by the parent and external_lhs while rhs is owned only + // by the parent. The getter adds one temporary owner to each field, so lhs must map and rhs may + // mutate in place. + ObjectRef external_lhs = + MakeMapLeaf(10, change_lhs ? TransformAction::kChange : TransformAction::kKeep); + ObjectRef parent_lhs = external_lhs; + ObjectRef rhs = MakeMapLeaf(20, change_rhs ? TransformAction::kChange : TransformAction::kKeep); + const Object* lhs_addr = external_lhs.get(); + const Object* rhs_addr = rhs.get(); + TPair root(std::move(parent_lhs), std::move(rhs)); + const Object* root_addr = root.get(); + ASSERT_EQ(root.use_count(), 1); + ASSERT_EQ(external_lhs.use_count(), 2); + ASSERT_EQ(root->rhs.use_count(), 1); + Any input = std::move(root); + TestMapper mapper = MakeTestMapper(); + + Expected result = mapper->InplaceMutateExpected(std::move(input)); + + ASSERT_TRUE(result.is_ok()); + TPair mapped = std::move(result).value().cast(); + EXPECT_EQ(mapped.get(), root_addr); + EXPECT_EQ(mapped.use_count(), 1); + + if (change_lhs) { + EXPECT_NE(mapped->lhs.get(), lhs_addr); + EXPECT_EQ(MapLeafValue(external_lhs), 10); + EXPECT_EQ(external_lhs.use_count(), 1); + EXPECT_EQ(MapLeafValue(mapped->lhs), 11); + EXPECT_EQ(mapped->lhs.use_count(), 1); + } else { + EXPECT_EQ(mapped->lhs.get(), lhs_addr); + EXPECT_EQ(external_lhs.use_count(), 2); + EXPECT_EQ(MapLeafValue(mapped->lhs), 10); + } + + EXPECT_EQ(mapped->rhs.get(), rhs_addr); + EXPECT_EQ(MapLeafValue(mapped->rhs), change_rhs ? 21 : 20); + EXPECT_EQ(mapped->rhs.use_count(), 1); + ExpectTrace({"map:10", "inplace:20"}); +} + +TEST(StructuralMapper, InplaceMutateSelectsEachFieldFromLogicalUniqueness) { + CheckInplaceFieldSelection(false, false); + CheckInplaceFieldSelection(true, false); + CheckInplaceFieldSelection(false, true); + CheckInplaceFieldSelection(true, true); +} + +TEST(StructuralMapper, DispatchesFFIFunctionHooks) { + TestMapper mapper = MakeTestMapper(); + + ObjectRef unchanged_source = MakeFunctionLeaf(5); + const Object* unchanged_addr = unchanged_source.get(); + Any unchanged_input = std::move(unchanged_source); + Expected unchanged_result = mapper->MapExpected(std::move(unchanged_input)); + + ASSERT_TRUE(unchanged_result.is_ok()); + ObjectRef unchanged_mapped = std::move(unchanged_result).value().cast(); + EXPECT_EQ(unchanged_mapped.get(), unchanged_addr); + EXPECT_EQ(unchanged_mapped.use_count(), 1); + ExpectTrace({"function-map:5"}); + + ClearMapperTrace(); + ObjectRef source = MakeFunctionLeaf(10, TransformAction::kChange); + const Object* source_addr = source.get(); + + Expected map_result = mapper->MapExpected(source); + + ASSERT_TRUE(map_result.is_ok()); + ObjectRef mapped = std::move(map_result).value().cast(); + EXPECT_NE(mapped.get(), source_addr); + EXPECT_EQ(MapLeafValue(source), 10); + EXPECT_EQ(MapLeafValue(mapped), 11); + EXPECT_EQ(source.use_count(), 1); + EXPECT_EQ(mapped.use_count(), 1); + ExpectTrace({"function-map:10"}); + + ClearMapperTrace(); + ObjectRef inplace_source = MakeFunctionLeaf(20, TransformAction::kChange); + const Object* inplace_addr = inplace_source.get(); + Any input = std::move(inplace_source); + Expected inplace_result = mapper->InplaceMutateExpected(std::move(input)); + + ASSERT_TRUE(inplace_result.is_ok()); + ObjectRef inplace_mapped = std::move(inplace_result).value().cast(); + EXPECT_EQ(inplace_mapped.get(), inplace_addr); + EXPECT_EQ(MapLeafValue(inplace_mapped), 21); + EXPECT_EQ(inplace_mapped.use_count(), 1); + ExpectTrace({"function-inplace:20"}); +} + +TEST(StructuralMapper, RecordsMapErrorContext) { + ObjectRef failing_leaf = MakeMapLeaf(1, TransformAction::kError); + TVar rhs("rhs"); + TPair root(std::move(failing_leaf), std::move(rhs)); + const Object* root_addr = root.get(); + const Object* leaf_addr = root->lhs.get(); + TestMapper mapper = MakeTestMapper(); + + Expected result = mapper->MapExpected(root); + + ASSERT_TRUE(result.is_err()); + Error error = result.error(); + EXPECT_EQ(error.kind(), "ValueError"); + EXPECT_EQ(error.message(), "structural map leaf failed"); + Optional context = VisitErrorContext::TryGetFromError(error); + ASSERT_TRUE(context.has_value()); + const List& chain = context.value()->reverse_visit_pattern; + ASSERT_EQ(chain.size(), 2U); + EXPECT_EQ(chain[0].get(), leaf_addr); + EXPECT_EQ(chain[1].get(), root_addr); + ExpectTrace({"map:1"}); +} + +TEST(StructuralMapper, CombinedErrorContextDoesNotDuplicateRoot) { + ObjectRef failing_leaf = MakeMapLeaf(1, TransformAction::kError); + TVar rhs("rhs"); + TPair root(std::move(failing_leaf), std::move(rhs)); + const Object* root_addr = root.get(); + const Object* leaf_addr = root->lhs.get(); + Any input = std::move(root); + TestMapper mapper = MakeTestMapper(); + + Expected result = mapper->MapOrInplaceMutateExpected(std::move(input)); + + ASSERT_TRUE(result.is_err()); + Error error = result.error(); + Optional context = VisitErrorContext::TryGetFromError(error); + ASSERT_TRUE(context.has_value()); + const List& chain = context.value()->reverse_visit_pattern; + ASSERT_EQ(chain.size(), 2U); + EXPECT_EQ(chain[0].get(), leaf_addr); + EXPECT_EQ(chain[1].get(), root_addr); + ExpectTrace({"inplace:1"}); +} + +TEST(StructuralMapper, InplaceMutationIsNotRolledBackOnError) { + ObjectRef lhs = MakeMapLeaf(1, TransformAction::kChange); + ObjectRef rhs = MakeMapLeaf(2, TransformAction::kError); + TPair root(std::move(lhs), std::move(rhs)); + const Object* root_addr = root.get(); + const Object* lhs_addr = root->lhs.get(); + const Object* rhs_addr = root->rhs.get(); + Any input = std::move(root); + TestMapper mapper = MakeTestMapper(); + + Expected result = mapper->InplaceMutateExpected(std::move(input)); + + ASSERT_TRUE(result.is_err()); + Error error = result.error(); + Optional context = VisitErrorContext::TryGetFromError(error); + ASSERT_TRUE(context.has_value()); + const List& chain = context.value()->reverse_visit_pattern; + ASSERT_EQ(chain.size(), 2U); + EXPECT_EQ(chain[0].get(), rhs_addr); + EXPECT_EQ(chain[1].get(), root_addr); + + const TPairObj* partially_mutated = chain[1].as(); + ASSERT_NE(partially_mutated, nullptr); + EXPECT_EQ(partially_mutated->lhs.get(), lhs_addr); + EXPECT_EQ(MapLeafValue(partially_mutated->lhs), 2); + ExpectTrace({"inplace:1", "inplace:2"}); +} + +} // namespace