tqchen commented on code in PR #601: URL: https://github.com/apache/tvm-ffi/pull/601#discussion_r3326884979
########## include/tvm/ffi/extra/structural_visit.h: ########## @@ -0,0 +1,535 @@ +/* + * 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_visit.h + * \brief Structural visit implementation + */ +#ifndef TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_ +#define TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_ + +#include <tvm/ffi/any.h> +#include <tvm/ffi/c_api.h> +#include <tvm/ffi/container/variant.h> +#include <tvm/ffi/expected.h> +#include <tvm/ffi/function_details.h> +#include <tvm/ffi/optional.h> + +#include <cstddef> +#include <exception> +#include <optional> +#include <tuple> +#include <type_traits> +#include <utility> + +namespace tvm { +namespace ffi { + +/*! + * \brief Object node carrying the optional payload for an interrupted structural visit. + */ +class VisitInterruptObj : public Object { + public: + /*! \brief Payload returned with the interrupt, or FFI None for no payload. */ + Any value; + + VisitInterruptObj() = default; + /*! + * \brief Construct a VisitInterruptObj with a payload. + * \param value The payload carried by the interrupt. + */ + explicit VisitInterruptObj(Any value) : value(std::move(value)) {} + + /// \cond Doxygen_Suppress + static constexpr const int32_t _type_index = TypeIndex::kTVMFFIVisitInterrupt; + static const constexpr bool _type_final = true; + TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIVisitInterrupt, VisitInterruptObj, + Object); + /// \endcond +}; + +/*! + * \brief ObjectRef wrapper for VisitInterruptObj. + */ +class VisitInterrupt : public ObjectRef { + public: + /*! \brief Construct an interrupt with no payload. */ + VisitInterrupt() : VisitInterrupt(Any(nullptr)) {} + /*! + * \brief Construct an interrupt with a user-defined payload. + * \param value The payload carried by the interrupt. + */ + explicit VisitInterrupt(Any value) + : ObjectRef(make_object<VisitInterruptObj>(std::move(value))) {} + + /// \cond Doxygen_Suppress + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VisitInterrupt, ObjectRef, VisitInterruptObj); + /// \endcond +}; + +class StructuralVisitorObj; + +/*! + * \brief C++ vtable for structural visitor dispatch. + * + * The vtable gives \ref StructuralVisitorObj an explicit dispatch table instead + * of relying on C++ virtual methods. + */ +struct StructuralVisitorVTable { + /*! \brief Visit callback. */ + TVMFFIAny (*visit)(StructuralVisitorObj* visitor, const Object* value) noexcept = nullptr; Review Comment: document the argument and return value convention e.g. Error is passed through TVMFFIAny return result is moved ########## include/tvm/ffi/extra/structural_visit.h: ########## @@ -0,0 +1,535 @@ +/* + * 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_visit.h + * \brief Structural visit implementation + */ +#ifndef TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_ +#define TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_ + +#include <tvm/ffi/any.h> +#include <tvm/ffi/c_api.h> +#include <tvm/ffi/container/variant.h> +#include <tvm/ffi/expected.h> +#include <tvm/ffi/function_details.h> +#include <tvm/ffi/optional.h> + +#include <cstddef> +#include <exception> +#include <optional> +#include <tuple> +#include <type_traits> +#include <utility> + +namespace tvm { +namespace ffi { + +/*! + * \brief Object node carrying the optional payload for an interrupted structural visit. + */ +class VisitInterruptObj : public Object { + public: + /*! \brief Payload returned with the interrupt, or FFI None for no payload. */ + Any value; + + VisitInterruptObj() = default; + /*! + * \brief Construct a VisitInterruptObj with a payload. + * \param value The payload carried by the interrupt. + */ + explicit VisitInterruptObj(Any value) : value(std::move(value)) {} + + /// \cond Doxygen_Suppress + static constexpr const int32_t _type_index = TypeIndex::kTVMFFIVisitInterrupt; + static const constexpr bool _type_final = true; + TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIVisitInterrupt, VisitInterruptObj, + Object); + /// \endcond +}; + +/*! + * \brief ObjectRef wrapper for VisitInterruptObj. + */ +class VisitInterrupt : public ObjectRef { + public: + /*! \brief Construct an interrupt with no payload. */ + VisitInterrupt() : VisitInterrupt(Any(nullptr)) {} + /*! + * \brief Construct an interrupt with a user-defined payload. + * \param value The payload carried by the interrupt. + */ + explicit VisitInterrupt(Any value) + : ObjectRef(make_object<VisitInterruptObj>(std::move(value))) {} + + /// \cond Doxygen_Suppress + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VisitInterrupt, ObjectRef, VisitInterruptObj); + /// \endcond +}; + +class StructuralVisitorObj; + +/*! + * \brief C++ vtable for structural visitor dispatch. + * + * The vtable gives \ref StructuralVisitorObj an explicit dispatch table instead + * of relying on C++ virtual methods. + */ +struct StructuralVisitorVTable { + /*! \brief Visit callback. */ + TVMFFIAny (*visit)(StructuralVisitorObj* visitor, const Object* value) noexcept = nullptr; +}; + +/*! + * \brief Object node of a structural visitor. + * + * A structural visitor is an active traversal context. It carries the dispatch + * table used to visit each object and the current def-region state used by + * structural equality/hash semantics. The visitor is ref-counted so it can + * cross FFI boundaries, but one underlying visitor object should not be shared + * by overlapping top-level traversals. + */ +class StructuralVisitorObj : public Object { + public: + /*! \brief Construct the default structural visitor. */ + StructuralVisitorObj() : StructuralVisitorObj(VTable()) {} + + /*! + * \brief Visit a value, dispatching through this visitor's vtable. + * + * \param value The object to visit. + * \return ``std::nullopt`` to continue traversal, or a \ref VisitInterrupt + * to halt the entire visit. + */ + TVM_FFI_INLINE Optional<VisitInterrupt> Visit(const ObjectRef& value) { + return VisitExpected(value).value(); + } + + /*! + * \brief Visit a value, propagating error through expected return. + * + * \param value The object to visit. + * \return Expected interrupt state. An error means traversal failed. + */ + TVM_FFI_INLINE Expected<Optional<VisitInterrupt>> VisitExpected(const ObjectRef& value) noexcept { + return details::ExpectedUnsafe::MoveFromTVMFFIAny<Optional<VisitInterrupt>>( + (*vtable_->visit)(this, value.get())); + } + + /*! + * \brief Default visit behavior: type-attr lookup, then reflection fallback. + * + * The default path first checks the structural-visit type attribute for a + * type-specific override. If none is registered, it visits all reflected + * fields that participate in structural equality/hash. + */ + TVM_FFI_INLINE Optional<VisitInterrupt> DefaultVisit(const ObjectRef& value) { + return DefaultVisitExpected(value).value(); + } + + /*! + * \brief Default visit behavior, propagating error through expected return. + * + * \param value The object to visit. + * \return Expected interrupt state. An error means traversal failed. + */ + TVM_FFI_DLL Expected<Optional<VisitInterrupt>> DefaultVisitExpected( + const ObjectRef& value) noexcept; + + /*! + * \brief Temporarily switch the def-region context while invoking \p callback. + * + * This helper scopes updates to the traversal state used by def/use-region + * aware visitors. The previous state is restored when the callback returns + * or throws. + * + * \param kind The def-region kind to set during the callback. + * \param callback A nullary callable that performs recursive visiting. + * \return The value returned by \p callback. + */ + template <typename Callback> + TVM_FFI_INLINE auto WithDefRegionKind(TVMFFIDefRegionKind kind, Callback&& callback) { + class Scope { + public: + Scope(StructuralVisitorObj* visitor, TVMFFIDefRegionKind kind) + : visitor_(visitor), old_kind_(visitor->def_region_mode_) { + visitor_->def_region_mode_ = kind; + } + ~Scope() { visitor_->def_region_mode_ = old_kind_; } + Scope(const Scope&) = delete; + Scope& operator=(const Scope&) = delete; + + private: + StructuralVisitorObj* visitor_; + TVMFFIDefRegionKind old_kind_; + }; + Scope scope(this, kind); + return std::forward<Callback>(callback)(); + } + + /// \cond Doxygen_Suppress + static constexpr const bool _type_mutable = true; + TVM_FFI_DECLARE_OBJECT_INFO("ffi.StructuralVisitor", StructuralVisitorObj, Object); + /// \endcond + + protected: + /*! + * \brief Construct a structural visitor subclass with a custom dispatch vtable. + * + * \param vtable The non-null dispatch table for this visitor. + * + * \note This constructor is for internal subclasses. The vtable and its + * ``visit`` callback must be valid for the lifetime of the visitor. + */ + explicit StructuralVisitorObj(const StructuralVisitorVTable* vtable) : vtable_(vtable) {} + + /*! + * \brief Required ABI dispatch table. \ref StructuralVisitorVTable + * It must never be null on a constructed visitor. + */ + const StructuralVisitorVTable* vtable_ = nullptr; + + /*! + * \brief Current def-region context for structural equality/hash semantics. + * + * This is shared mutable traversal state. Be careful when mutating it through + * multiple references to the same visitor object. Use \ref WithDefRegionKind + * to scope temporary changes. + */ + TVMFFIDefRegionKind def_region_mode_ = kTVMFFIDefRegionKindNone; + + private: + /*! + * \brief Return the vtable used by the default visitor implementation. + * \return Pointer to the static structural visitor vtable. + */ + static const StructuralVisitorVTable* VTable() { + static const StructuralVisitorVTable vtable{&StructuralVisitorObj::DispatchVisit}; + return &vtable; + } + + /*! + * \brief Dispatch from the vtable to the default visitor implementation. + * \param visitor The structural visitor object. + * \param value The object to visit. + * \return Interrupt state, or an error if traversal failed. + */ + static TVMFFIAny DispatchVisit(StructuralVisitorObj* visitor, const Object* value) noexcept { + ObjectRef value_ref( + details::ObjectUnsafe::ObjectPtrFromUnowned<Object>(const_cast<Object*>(value))); + return details::ExpectedUnsafe::MoveToTVMFFIAny(visitor->DefaultVisitExpected(value_ref)); + } +}; + +/*! + * \brief ObjectRef wrapper of \ref StructuralVisitorObj. + * + * \sa StructuralVisitorObj + */ +class StructuralVisitor : public ObjectRef { + public: + /*! + * \brief Construct the default structural visitor. + */ + StructuralVisitor() : ObjectRef(make_object<StructuralVisitorObj>()) {} + /*! + * \brief Construct from an existing object pointer. + * \param n The object pointer to wrap. + */ + explicit StructuralVisitor(ObjectPtr<StructuralVisitorObj> n) : ObjectRef(std::move(n)) {} + + /// \cond Doxygen_Suppress + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StructuralVisitor, ObjectRef, StructuralVisitorObj); + /// \endcond +}; + +// --------------------------------------------------------------------------- +// Structural Walk API. +// --------------------------------------------------------------------------- + +/*! + * \brief Callback order for \ref StructuralWalk. + */ +enum class WalkOrder : int32_t { + /*! \brief Invoke the callback before visiting children. */ + kPreOrder = 0, + /*! \brief Invoke the callback after visiting children. */ + kPostOrder = 1, +}; + +namespace details { + +/*! + * \brief Return true when \p result already carries a traversal-stopping state. + * \tparam T The Expected success type. + * \param result The Expected value to inspect. + * \return Whether \p result stores an Error or VisitInterrupt. + */ +template <typename T> +TVM_FFI_INLINE bool NeedEarlyReturn(const Expected<T>& result) noexcept { + int32_t type_index = result.type_index(); + return type_index == TypeIndex::kTVMFFIError || type_index == TypeIndex::kTVMFFIVisitInterrupt; +} + +/*! + * \brief Return from the current ABI visit function if \p Result stops traversal. + * + * \param Result Expression returning an Expected value whose raw storage can be + * moved to TVMFFIAny. + */ +#define TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Result) \ + if (auto result = (Result); \ + TVM_FFI_PREDICT_FALSE(::tvm::ffi::details::NeedEarlyReturn(result))) { \ + return ::tvm::ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result)); \ + } + +/*! + * \brief Move raw Expected storage into another Expected type. + * \tparam T The destination Expected success type. + * \tparam U The source Expected success type. + * \param result The source Expected value to move from. + * \return Expected<T> backed by the moved raw storage. + * + * \note This assumes the stored value is ABI-compatible with T, or is Error. + */ +template <typename T, typename U> +TVM_FFI_INLINE Expected<T> UnsafeMoveExpected(Expected<U>&& result) { Review Comment: this needs to be inlined, since there is already unsafe, or move this into ExpectedUnsafe ########## include/tvm/ffi/extra/structural_visit.h: ########## @@ -0,0 +1,535 @@ +/* + * 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_visit.h + * \brief Structural visit implementation + */ +#ifndef TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_ +#define TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_ + +#include <tvm/ffi/any.h> +#include <tvm/ffi/c_api.h> +#include <tvm/ffi/container/variant.h> +#include <tvm/ffi/expected.h> +#include <tvm/ffi/function_details.h> +#include <tvm/ffi/optional.h> + +#include <cstddef> +#include <exception> +#include <optional> +#include <tuple> +#include <type_traits> +#include <utility> + +namespace tvm { +namespace ffi { + +/*! + * \brief Object node carrying the optional payload for an interrupted structural visit. + */ +class VisitInterruptObj : public Object { + public: + /*! \brief Payload returned with the interrupt, or FFI None for no payload. */ + Any value; + + VisitInterruptObj() = default; + /*! + * \brief Construct a VisitInterruptObj with a payload. + * \param value The payload carried by the interrupt. + */ + explicit VisitInterruptObj(Any value) : value(std::move(value)) {} + + /// \cond Doxygen_Suppress + static constexpr const int32_t _type_index = TypeIndex::kTVMFFIVisitInterrupt; + static const constexpr bool _type_final = true; + TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIVisitInterrupt, VisitInterruptObj, + Object); + /// \endcond +}; + +/*! + * \brief ObjectRef wrapper for VisitInterruptObj. + */ +class VisitInterrupt : public ObjectRef { + public: + /*! \brief Construct an interrupt with no payload. */ + VisitInterrupt() : VisitInterrupt(Any(nullptr)) {} + /*! + * \brief Construct an interrupt with a user-defined payload. + * \param value The payload carried by the interrupt. + */ + explicit VisitInterrupt(Any value) + : ObjectRef(make_object<VisitInterruptObj>(std::move(value))) {} + + /// \cond Doxygen_Suppress + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VisitInterrupt, ObjectRef, VisitInterruptObj); + /// \endcond +}; + +class StructuralVisitorObj; + +/*! + * \brief C++ vtable for structural visitor dispatch. + * + * The vtable gives \ref StructuralVisitorObj an explicit dispatch table instead + * of relying on C++ virtual methods. + */ +struct StructuralVisitorVTable { + /*! \brief Visit callback. */ + TVMFFIAny (*visit)(StructuralVisitorObj* visitor, const Object* value) noexcept = nullptr; +}; + +/*! + * \brief Object node of a structural visitor. + * + * A structural visitor is an active traversal context. It carries the dispatch + * table used to visit each object and the current def-region state used by + * structural equality/hash semantics. The visitor is ref-counted so it can + * cross FFI boundaries, but one underlying visitor object should not be shared + * by overlapping top-level traversals. + */ +class StructuralVisitorObj : public Object { + public: + /*! \brief Construct the default structural visitor. */ + StructuralVisitorObj() : StructuralVisitorObj(VTable()) {} + + /*! + * \brief Visit a value, dispatching through this visitor's vtable. + * + * \param value The object to visit. + * \return ``std::nullopt`` to continue traversal, or a \ref VisitInterrupt + * to halt the entire visit. + */ + TVM_FFI_INLINE Optional<VisitInterrupt> Visit(const ObjectRef& value) { + return VisitExpected(value).value(); + } + + /*! + * \brief Visit a value, propagating error through expected return. + * + * \param value The object to visit. + * \return Expected interrupt state. An error means traversal failed. + */ + TVM_FFI_INLINE Expected<Optional<VisitInterrupt>> VisitExpected(const ObjectRef& value) noexcept { + return details::ExpectedUnsafe::MoveFromTVMFFIAny<Optional<VisitInterrupt>>( + (*vtable_->visit)(this, value.get())); + } + + /*! + * \brief Default visit behavior: type-attr lookup, then reflection fallback. + * + * The default path first checks the structural-visit type attribute for a + * type-specific override. If none is registered, it visits all reflected + * fields that participate in structural equality/hash. + */ + TVM_FFI_INLINE Optional<VisitInterrupt> DefaultVisit(const ObjectRef& value) { + return DefaultVisitExpected(value).value(); + } + + /*! + * \brief Default visit behavior, propagating error through expected return. + * + * \param value The object to visit. + * \return Expected interrupt state. An error means traversal failed. + */ + TVM_FFI_DLL Expected<Optional<VisitInterrupt>> DefaultVisitExpected( + const ObjectRef& value) noexcept; + + /*! + * \brief Temporarily switch the def-region context while invoking \p callback. + * + * This helper scopes updates to the traversal state used by def/use-region + * aware visitors. The previous state is restored when the callback returns + * or throws. + * + * \param kind The def-region kind to set during the callback. + * \param callback A nullary callable that performs recursive visiting. + * \return The value returned by \p callback. + */ + template <typename Callback> + TVM_FFI_INLINE auto WithDefRegionKind(TVMFFIDefRegionKind kind, Callback&& callback) { + class Scope { + public: + Scope(StructuralVisitorObj* visitor, TVMFFIDefRegionKind kind) + : visitor_(visitor), old_kind_(visitor->def_region_mode_) { + visitor_->def_region_mode_ = kind; + } + ~Scope() { visitor_->def_region_mode_ = old_kind_; } + Scope(const Scope&) = delete; + Scope& operator=(const Scope&) = delete; + + private: + StructuralVisitorObj* visitor_; + TVMFFIDefRegionKind old_kind_; + }; + Scope scope(this, kind); + return std::forward<Callback>(callback)(); + } + + /// \cond Doxygen_Suppress + static constexpr const bool _type_mutable = true; + TVM_FFI_DECLARE_OBJECT_INFO("ffi.StructuralVisitor", StructuralVisitorObj, Object); + /// \endcond + + protected: + /*! + * \brief Construct a structural visitor subclass with a custom dispatch vtable. + * + * \param vtable The non-null dispatch table for this visitor. + * + * \note This constructor is for internal subclasses. The vtable and its + * ``visit`` callback must be valid for the lifetime of the visitor. + */ + explicit StructuralVisitorObj(const StructuralVisitorVTable* vtable) : vtable_(vtable) {} + + /*! + * \brief Required ABI dispatch table. \ref StructuralVisitorVTable + * It must never be null on a constructed visitor. + */ + const StructuralVisitorVTable* vtable_ = nullptr; + + /*! + * \brief Current def-region context for structural equality/hash semantics. + * + * This is shared mutable traversal state. Be careful when mutating it through + * multiple references to the same visitor object. Use \ref WithDefRegionKind + * to scope temporary changes. + */ + TVMFFIDefRegionKind def_region_mode_ = kTVMFFIDefRegionKindNone; + + private: + /*! + * \brief Return the vtable used by the default visitor implementation. + * \return Pointer to the static structural visitor vtable. + */ + static const StructuralVisitorVTable* VTable() { + static const StructuralVisitorVTable vtable{&StructuralVisitorObj::DispatchVisit}; + return &vtable; + } + + /*! + * \brief Dispatch from the vtable to the default visitor implementation. + * \param visitor The structural visitor object. + * \param value The object to visit. + * \return Interrupt state, or an error if traversal failed. + */ + static TVMFFIAny DispatchVisit(StructuralVisitorObj* visitor, const Object* value) noexcept { + ObjectRef value_ref( + details::ObjectUnsafe::ObjectPtrFromUnowned<Object>(const_cast<Object*>(value))); + return details::ExpectedUnsafe::MoveToTVMFFIAny(visitor->DefaultVisitExpected(value_ref)); + } +}; + +/*! + * \brief ObjectRef wrapper of \ref StructuralVisitorObj. + * + * \sa StructuralVisitorObj + */ +class StructuralVisitor : public ObjectRef { + public: + /*! + * \brief Construct the default structural visitor. + */ + StructuralVisitor() : ObjectRef(make_object<StructuralVisitorObj>()) {} + /*! + * \brief Construct from an existing object pointer. + * \param n The object pointer to wrap. + */ + explicit StructuralVisitor(ObjectPtr<StructuralVisitorObj> n) : ObjectRef(std::move(n)) {} + + /// \cond Doxygen_Suppress + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StructuralVisitor, ObjectRef, StructuralVisitorObj); + /// \endcond +}; + +// --------------------------------------------------------------------------- +// Structural Walk API. +// --------------------------------------------------------------------------- + +/*! + * \brief Callback order for \ref StructuralWalk. + */ +enum class WalkOrder : int32_t { + /*! \brief Invoke the callback before visiting children. */ + kPreOrder = 0, + /*! \brief Invoke the callback after visiting children. */ + kPostOrder = 1, +}; + +namespace details { + +/*! + * \brief Return true when \p result already carries a traversal-stopping state. + * \tparam T The Expected success type. + * \param result The Expected value to inspect. + * \return Whether \p result stores an Error or VisitInterrupt. + */ +template <typename T> +TVM_FFI_INLINE bool NeedEarlyReturn(const Expected<T>& result) noexcept { Review Comment: naming in public header needs more clarity, StructuralVisitNeedEarlyReturn ########## include/tvm/ffi/extra/structural_visit.h: ########## @@ -0,0 +1,535 @@ +/* + * 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_visit.h + * \brief Structural visit implementation + */ +#ifndef TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_ +#define TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_ + +#include <tvm/ffi/any.h> +#include <tvm/ffi/c_api.h> +#include <tvm/ffi/container/variant.h> +#include <tvm/ffi/expected.h> +#include <tvm/ffi/function_details.h> +#include <tvm/ffi/optional.h> + +#include <cstddef> +#include <exception> +#include <optional> +#include <tuple> +#include <type_traits> +#include <utility> + +namespace tvm { +namespace ffi { + +/*! + * \brief Object node carrying the optional payload for an interrupted structural visit. + */ +class VisitInterruptObj : public Object { + public: + /*! \brief Payload returned with the interrupt, or FFI None for no payload. */ + Any value; + + VisitInterruptObj() = default; + /*! + * \brief Construct a VisitInterruptObj with a payload. + * \param value The payload carried by the interrupt. + */ + explicit VisitInterruptObj(Any value) : value(std::move(value)) {} + + /// \cond Doxygen_Suppress + static constexpr const int32_t _type_index = TypeIndex::kTVMFFIVisitInterrupt; + static const constexpr bool _type_final = true; + TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIVisitInterrupt, VisitInterruptObj, + Object); + /// \endcond +}; + +/*! + * \brief ObjectRef wrapper for VisitInterruptObj. + */ +class VisitInterrupt : public ObjectRef { + public: + /*! \brief Construct an interrupt with no payload. */ + VisitInterrupt() : VisitInterrupt(Any(nullptr)) {} + /*! + * \brief Construct an interrupt with a user-defined payload. + * \param value The payload carried by the interrupt. + */ + explicit VisitInterrupt(Any value) + : ObjectRef(make_object<VisitInterruptObj>(std::move(value))) {} + + /// \cond Doxygen_Suppress + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VisitInterrupt, ObjectRef, VisitInterruptObj); + /// \endcond +}; + +class StructuralVisitorObj; + +/*! + * \brief C++ vtable for structural visitor dispatch. + * + * The vtable gives \ref StructuralVisitorObj an explicit dispatch table instead + * of relying on C++ virtual methods. + */ +struct StructuralVisitorVTable { + /*! \brief Visit callback. */ + TVMFFIAny (*visit)(StructuralVisitorObj* visitor, const Object* value) noexcept = nullptr; +}; + +/*! + * \brief Object node of a structural visitor. + * + * A structural visitor is an active traversal context. It carries the dispatch + * table used to visit each object and the current def-region state used by + * structural equality/hash semantics. The visitor is ref-counted so it can + * cross FFI boundaries, but one underlying visitor object should not be shared + * by overlapping top-level traversals. + */ +class StructuralVisitorObj : public Object { + public: + /*! \brief Construct the default structural visitor. */ + StructuralVisitorObj() : StructuralVisitorObj(VTable()) {} + + /*! + * \brief Visit a value, dispatching through this visitor's vtable. + * + * \param value The object to visit. + * \return ``std::nullopt`` to continue traversal, or a \ref VisitInterrupt + * to halt the entire visit. + */ + TVM_FFI_INLINE Optional<VisitInterrupt> Visit(const ObjectRef& value) { + return VisitExpected(value).value(); + } + + /*! + * \brief Visit a value, propagating error through expected return. + * + * \param value The object to visit. + * \return Expected interrupt state. An error means traversal failed. + */ + TVM_FFI_INLINE Expected<Optional<VisitInterrupt>> VisitExpected(const ObjectRef& value) noexcept { + return details::ExpectedUnsafe::MoveFromTVMFFIAny<Optional<VisitInterrupt>>( + (*vtable_->visit)(this, value.get())); + } + + /*! + * \brief Default visit behavior: type-attr lookup, then reflection fallback. + * + * The default path first checks the structural-visit type attribute for a Review Comment: Impelmentation note should not be part of the documentation, just say the DefaultVisit based on registered type attr of each type ########## include/tvm/ffi/expected.h: ########## @@ -186,9 +195,106 @@ class Expected { } private: + Expected() = default; + + friend struct details::ExpectedUnsafe; + Any data_; // Invariant: holds a T (type_index != kTVMFFIError) or an Error. }; +namespace details { + +/*! + * \brief Unsafe raw-storage helpers for Expected. + * + * These helpers bypass normal value checking and are intended for ABI boundaries + * that already know the underlying Any storage holds either a valid T or Error. + */ +struct ExpectedUnsafe { + /*! + * \brief Copy a raw TVMFFIAny into Expected storage. + * \tparam T The Expected success type. + * \param raw The raw FFI value to copy. + * \return Expected backed by a copied Any storage. + */ + template <typename T> + TVM_FFI_INLINE static Expected<T> CopyFromTVMFFIAny(TVMFFIAny raw) { + Expected<T> result; + result.data_ = Any(AnyView::CopyFromTVMFFIAny(raw)); + return result; + } + + /*! + * \brief Move a raw TVMFFIAny into Expected storage. + * \tparam T The Expected success type. + * \param raw The raw FFI value to move. + * \return Expected backed by moved Any storage. + */ + template <typename T> + TVM_FFI_INLINE static Expected<T> MoveFromTVMFFIAny(TVMFFIAny raw) { + Expected<T> result; + result.data_ = AnyUnsafe::MoveTVMFFIAnyToAny(&raw); + return result; + } + + /*! + * \brief Copy Expected storage to a raw TVMFFIAny. + * \tparam T The Expected success type. + * \param result The Expected value to copy from. + * \return Raw FFI value containing a copy of the underlying Any storage. + */ + template <typename T> + TVM_FFI_INLINE static TVMFFIAny CopyToTVMFFIAny(const Expected<T>& result) { + return AnyUnsafe::MoveAnyToTVMFFIAny(Any(result.data_)); + } + + /*! + * \brief Move Expected storage to a raw TVMFFIAny. + * \tparam T The Expected success type. + * \param result The Expected value to move from. + * \return Raw FFI value containing moved underlying Any storage. + */ + template <typename T> + TVM_FFI_INLINE static TVMFFIAny MoveToTVMFFIAny(Expected<T>&& result) { + return AnyUnsafe::MoveAnyToTVMFFIAny(std::move(result.data_)); + } + + /*! + * \brief Return the underlying Any storage. + * \tparam T The Expected success type. + * \param result The Expected value to inspect. + * \return Const reference to the raw Any storage. + */ + template <typename T> + TVM_FFI_INLINE static const Any& GetData(const Expected<T>& result) noexcept { + return result.data_; + } + + /*! + * \brief Return mutable underlying Any storage. + * \tparam T The Expected success type. + * \param result The Expected value to inspect. + * \return Mutable reference to the raw Any storage. + */ + template <typename T> + TVM_FFI_INLINE static Any& GetData(Expected<T>& result) noexcept { Review Comment: only expose what we need we only need GetData fthat returns const ref ########## include/tvm/ffi/expected.h: ########## @@ -186,9 +195,106 @@ class Expected { } private: + Expected() = default; + + friend struct details::ExpectedUnsafe; + Any data_; // Invariant: holds a T (type_index != kTVMFFIError) or an Error. }; +namespace details { + +/*! + * \brief Unsafe raw-storage helpers for Expected. + * + * These helpers bypass normal value checking and are intended for ABI boundaries + * that already know the underlying Any storage holds either a valid T or Error. + */ +struct ExpectedUnsafe { + /*! + * \brief Copy a raw TVMFFIAny into Expected storage. + * \tparam T The Expected success type. + * \param raw The raw FFI value to copy. + * \return Expected backed by a copied Any storage. + */ + template <typename T> + TVM_FFI_INLINE static Expected<T> CopyFromTVMFFIAny(TVMFFIAny raw) { + Expected<T> result; Review Comment: This is not used, don't expose ########## include/tvm/ffi/extra/structural_visit.h: ########## @@ -0,0 +1,535 @@ +/* + * 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_visit.h + * \brief Structural visit implementation + */ +#ifndef TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_ +#define TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_ + +#include <tvm/ffi/any.h> +#include <tvm/ffi/c_api.h> +#include <tvm/ffi/container/variant.h> +#include <tvm/ffi/expected.h> +#include <tvm/ffi/function_details.h> +#include <tvm/ffi/optional.h> + +#include <cstddef> +#include <exception> +#include <optional> +#include <tuple> +#include <type_traits> +#include <utility> + +namespace tvm { +namespace ffi { + +/*! + * \brief Object node carrying the optional payload for an interrupted structural visit. + */ +class VisitInterruptObj : public Object { + public: + /*! \brief Payload returned with the interrupt, or FFI None for no payload. */ + Any value; + + VisitInterruptObj() = default; + /*! + * \brief Construct a VisitInterruptObj with a payload. + * \param value The payload carried by the interrupt. + */ + explicit VisitInterruptObj(Any value) : value(std::move(value)) {} + + /// \cond Doxygen_Suppress + static constexpr const int32_t _type_index = TypeIndex::kTVMFFIVisitInterrupt; + static const constexpr bool _type_final = true; + TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIVisitInterrupt, VisitInterruptObj, + Object); + /// \endcond +}; + +/*! + * \brief ObjectRef wrapper for VisitInterruptObj. + */ +class VisitInterrupt : public ObjectRef { + public: + /*! \brief Construct an interrupt with no payload. */ + VisitInterrupt() : VisitInterrupt(Any(nullptr)) {} + /*! + * \brief Construct an interrupt with a user-defined payload. + * \param value The payload carried by the interrupt. + */ + explicit VisitInterrupt(Any value) + : ObjectRef(make_object<VisitInterruptObj>(std::move(value))) {} + + /// \cond Doxygen_Suppress + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VisitInterrupt, ObjectRef, VisitInterruptObj); + /// \endcond +}; + +class StructuralVisitorObj; + +/*! + * \brief C++ vtable for structural visitor dispatch. + * + * The vtable gives \ref StructuralVisitorObj an explicit dispatch table instead + * of relying on C++ virtual methods. + */ +struct StructuralVisitorVTable { + /*! \brief Visit callback. */ + TVMFFIAny (*visit)(StructuralVisitorObj* visitor, const Object* value) noexcept = nullptr; +}; + +/*! + * \brief Object node of a structural visitor. + * + * A structural visitor is an active traversal context. It carries the dispatch + * table used to visit each object and the current def-region state used by + * structural equality/hash semantics. The visitor is ref-counted so it can + * cross FFI boundaries, but one underlying visitor object should not be shared + * by overlapping top-level traversals. + */ +class StructuralVisitorObj : public Object { + public: + /*! \brief Construct the default structural visitor. */ + StructuralVisitorObj() : StructuralVisitorObj(VTable()) {} + + /*! + * \brief Visit a value, dispatching through this visitor's vtable. + * + * \param value The object to visit. + * \return ``std::nullopt`` to continue traversal, or a \ref VisitInterrupt + * to halt the entire visit. + */ + TVM_FFI_INLINE Optional<VisitInterrupt> Visit(const ObjectRef& value) { + return VisitExpected(value).value(); + } + + /*! + * \brief Visit a value, propagating error through expected return. + * + * \param value The object to visit. + * \return Expected interrupt state. An error means traversal failed. + */ + TVM_FFI_INLINE Expected<Optional<VisitInterrupt>> VisitExpected(const ObjectRef& value) noexcept { + return details::ExpectedUnsafe::MoveFromTVMFFIAny<Optional<VisitInterrupt>>( + (*vtable_->visit)(this, value.get())); + } + + /*! + * \brief Default visit behavior: type-attr lookup, then reflection fallback. + * + * The default path first checks the structural-visit type attribute for a + * type-specific override. If none is registered, it visits all reflected + * fields that participate in structural equality/hash. + */ + TVM_FFI_INLINE Optional<VisitInterrupt> DefaultVisit(const ObjectRef& value) { + return DefaultVisitExpected(value).value(); + } + + /*! + * \brief Default visit behavior, propagating error through expected return. + * + * \param value The object to visit. + * \return Expected interrupt state. An error means traversal failed. + */ + TVM_FFI_DLL Expected<Optional<VisitInterrupt>> DefaultVisitExpected( + const ObjectRef& value) noexcept; + + /*! + * \brief Temporarily switch the def-region context while invoking \p callback. + * + * This helper scopes updates to the traversal state used by def/use-region + * aware visitors. The previous state is restored when the callback returns + * or throws. + * + * \param kind The def-region kind to set during the callback. + * \param callback A nullary callable that performs recursive visiting. + * \return The value returned by \p callback. + */ + template <typename Callback> + TVM_FFI_INLINE auto WithDefRegionKind(TVMFFIDefRegionKind kind, Callback&& callback) { + class Scope { + public: + Scope(StructuralVisitorObj* visitor, TVMFFIDefRegionKind kind) + : visitor_(visitor), old_kind_(visitor->def_region_mode_) { + visitor_->def_region_mode_ = kind; + } + ~Scope() { visitor_->def_region_mode_ = old_kind_; } + Scope(const Scope&) = delete; + Scope& operator=(const Scope&) = delete; + + private: + StructuralVisitorObj* visitor_; + TVMFFIDefRegionKind old_kind_; + }; + Scope scope(this, kind); + return std::forward<Callback>(callback)(); + } + + /// \cond Doxygen_Suppress + static constexpr const bool _type_mutable = true; + TVM_FFI_DECLARE_OBJECT_INFO("ffi.StructuralVisitor", StructuralVisitorObj, Object); + /// \endcond + + protected: + /*! + * \brief Construct a structural visitor subclass with a custom dispatch vtable. + * + * \param vtable The non-null dispatch table for this visitor. + * + * \note This constructor is for internal subclasses. The vtable and its + * ``visit`` callback must be valid for the lifetime of the visitor. + */ + explicit StructuralVisitorObj(const StructuralVisitorVTable* vtable) : vtable_(vtable) {} + + /*! + * \brief Required ABI dispatch table. \ref StructuralVisitorVTable + * It must never be null on a constructed visitor. + */ + const StructuralVisitorVTable* vtable_ = nullptr; + + /*! + * \brief Current def-region context for structural equality/hash semantics. + * + * This is shared mutable traversal state. Be careful when mutating it through + * multiple references to the same visitor object. Use \ref WithDefRegionKind + * to scope temporary changes. + */ + TVMFFIDefRegionKind def_region_mode_ = kTVMFFIDefRegionKindNone; + + private: + /*! + * \brief Return the vtable used by the default visitor implementation. + * \return Pointer to the static structural visitor vtable. + */ + static const StructuralVisitorVTable* VTable() { + static const StructuralVisitorVTable vtable{&StructuralVisitorObj::DispatchVisit}; + return &vtable; + } + + /*! + * \brief Dispatch from the vtable to the default visitor implementation. + * \param visitor The structural visitor object. + * \param value The object to visit. + * \return Interrupt state, or an error if traversal failed. + */ + static TVMFFIAny DispatchVisit(StructuralVisitorObj* visitor, const Object* value) noexcept { + ObjectRef value_ref( + details::ObjectUnsafe::ObjectPtrFromUnowned<Object>(const_cast<Object*>(value))); + return details::ExpectedUnsafe::MoveToTVMFFIAny(visitor->DefaultVisitExpected(value_ref)); + } +}; + +/*! + * \brief ObjectRef wrapper of \ref StructuralVisitorObj. + * + * \sa StructuralVisitorObj + */ +class StructuralVisitor : public ObjectRef { + public: + /*! + * \brief Construct the default structural visitor. + */ + StructuralVisitor() : ObjectRef(make_object<StructuralVisitorObj>()) {} + /*! + * \brief Construct from an existing object pointer. + * \param n The object pointer to wrap. + */ + explicit StructuralVisitor(ObjectPtr<StructuralVisitorObj> n) : ObjectRef(std::move(n)) {} + + /// \cond Doxygen_Suppress + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StructuralVisitor, ObjectRef, StructuralVisitorObj); + /// \endcond +}; + +// --------------------------------------------------------------------------- +// Structural Walk API. +// --------------------------------------------------------------------------- + +/*! + * \brief Callback order for \ref StructuralWalk. + */ +enum class WalkOrder : int32_t { + /*! \brief Invoke the callback before visiting children. */ + kPreOrder = 0, + /*! \brief Invoke the callback after visiting children. */ + kPostOrder = 1, +}; + +namespace details { + +/*! + * \brief Return true when \p result already carries a traversal-stopping state. + * \tparam T The Expected success type. + * \param result The Expected value to inspect. + * \return Whether \p result stores an Error or VisitInterrupt. + */ +template <typename T> +TVM_FFI_INLINE bool NeedEarlyReturn(const Expected<T>& result) noexcept { + int32_t type_index = result.type_index(); + return type_index == TypeIndex::kTVMFFIError || type_index == TypeIndex::kTVMFFIVisitInterrupt; +} + +/*! + * \brief Return from the current ABI visit function if \p Result stops traversal. + * + * \param Result Expression returning an Expected value whose raw storage can be + * moved to TVMFFIAny. + */ +#define TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Result) \ + if (auto result = (Result); \ + TVM_FFI_PREDICT_FALSE(::tvm::ffi::details::NeedEarlyReturn(result))) { \ + return ::tvm::ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result)); \ + } + +/*! + * \brief Move raw Expected storage into another Expected type. + * \tparam T The destination Expected success type. + * \tparam U The source Expected success type. + * \param result The source Expected value to move from. + * \return Expected<T> backed by the moved raw storage. + * + * \note This assumes the stored value is ABI-compatible with T, or is Error. + */ +template <typename T, typename U> +TVM_FFI_INLINE Expected<T> UnsafeMoveExpected(Expected<U>&& result) { + return ExpectedUnsafe::MoveFromTVMFFIAny<T>(ExpectedUnsafe::MoveToTVMFFIAny(std::move(result))); +} + +/*! + * \brief Read an Expected success value as a compatible raw storage type. + * \tparam T The type to read from the underlying Any storage. + * \tparam U The Expected success type. + * \param result The Expected value to read from. + * \return The stored value decoded as T. + * + * \note This assumes \p result stores T-compatible Any storage, or Error. + */ +template <typename T, typename U> +TVM_FFI_INLINE T UnsafeExpectedValueAs(const Expected<U>& result) { + const Any& data = ExpectedUnsafe::GetData(result); + if (TVM_FFI_PREDICT_TRUE(data.type_index() != TypeIndex::kTVMFFIError)) { + return AnyUnsafe::CopyFromAnyViewAfterCheck<T>(data); + } + throw AnyUnsafe::CopyFromAnyViewAfterCheck<Error>(data); +} + +/*! + * \brief Concrete visitor implementation used by callback-dispatched ``StructuralWalk``. + * \tparam order Callback placement relative to child traversal. + * \tparam Callbacks Callback types to invoke on matching nodes. + */ +template <WalkOrder order, typename... Callbacks> +class StructuralWalkCallbackVisitorObj : public StructuralVisitorObj { + public: + /*! + * \brief Construct a structural walk visitor. + * \param callbacks Callbacks to invoke on matching nodes. + */ + template <typename... F> + explicit StructuralWalkCallbackVisitorObj(F&&... callbacks) + : StructuralVisitorObj(VTable()), callbacks_(std::forward<F>(callbacks)...) {} + + private: + /*! + * \brief Return the vtable used by this visitor implementation. + * \return Pointer to the static structural visitor vtable. + */ + static const StructuralVisitorVTable* VTable() { + static const StructuralVisitorVTable vtable{&StructuralWalkCallbackVisitorObj::DispatchVisit}; + return &vtable; + } + + /*! + * \brief Dispatch from the erased visitor pointer to the concrete walk visitor. + * \param self The erased structural visitor object. + * \param value The object to visit. + * \return Interrupt state, or an error if traversal failed. + */ + static TVMFFIAny DispatchVisit(StructuralVisitorObj* self, const Object* value) noexcept { + ObjectRef value_ref( + details::ObjectUnsafe::ObjectPtrFromUnowned<Object>(const_cast<Object*>(value))); + return details::ExpectedUnsafe::MoveToTVMFFIAny( + static_cast<StructuralWalkCallbackVisitorObj*>(self)->VisitImpl(value_ref)); + } + + /*! + * \brief Visit one object according to the configured walk order. + * \param value The object to visit. + * \return Interrupt state, or an error if traversal failed. + */ + Expected<Optional<VisitInterrupt>> VisitImpl(const ObjectRef& value) noexcept { + try { + if constexpr (order == WalkOrder::kPreOrder) { + auto result = InvokeMatchingCallback(value); + if (TVM_FFI_PREDICT_FALSE(NeedEarlyReturn(result))) { + return UnsafeMoveExpected<Optional<VisitInterrupt>>(std::move(result)); + } + int32_t type_index = result.type_index(); + TVM_FFI_UNSAFE_ASSUME(type_index == TypeIndex::kTVMFFIBool); + if (TVM_FFI_PREDICT_FALSE(!UnsafeExpectedValueAs<bool>(result))) { + return Optional<VisitInterrupt>(std::nullopt); + } + } + + Expected<Optional<VisitInterrupt>> interrupt = DefaultVisitExpected(value); + if (TVM_FFI_PREDICT_FALSE(NeedEarlyReturn(interrupt))) { + return interrupt; + } + + if constexpr (order == WalkOrder::kPostOrder) { + auto result = InvokeMatchingCallback(value); + if (TVM_FFI_PREDICT_FALSE(NeedEarlyReturn(result))) { + return UnsafeMoveExpected<Optional<VisitInterrupt>>(std::move(result)); + } + } + + return Optional<VisitInterrupt>(std::nullopt); + } catch (const Error& err) { + return Unexpected(err); + } catch (const std::exception& err) { + return Unexpected(Error("InternalError", err.what(), "")); + } catch (...) { + return Unexpected(Error("InternalError", "Unknown structural walk error", "")); + } + } + + /*! + * \brief Invoke the first callback whose argument type matches ``value``. + * \param value The object to match against the callback arguments. + * \return Callback action, or an error if the callback failed. + */ + Expected<Variant<VisitInterrupt, bool>> InvokeMatchingCallback(const ObjectRef& value) { + if (TVM_FFI_PREDICT_FALSE(!value.defined())) { + return Variant<VisitInterrupt, bool>(true); + } + return InvokeMatchingCallback<0>(value); + } + + /*! + * \brief Invoke the first matching callback from index ``I`` onward. + * \param value The object to match. + * \return Callback action, or continue if no callback matched. + */ + template <std::size_t I> + Expected<Variant<VisitInterrupt, bool>> InvokeMatchingCallback(const ObjectRef& value) { + if constexpr (I == sizeof...(Callbacks)) { + return Variant<VisitInterrupt, bool>(true); + } else { + using Callback = std::tuple_element_t<I, std::tuple<Callbacks...>>; + using FuncInfo = FunctionInfo<Callback>; + using Arg = std::tuple_element_t<0, typename FuncInfo::ArgType>; + using RefType = std::remove_cv_t<std::remove_reference_t<Arg>>; + + if (auto ref = value.template as<RefType>()) { + return std::get<I>(callbacks_)(*ref); + } + return InvokeMatchingCallback<I + 1>(value); + } + } + + /*! \brief Concrete callback objects. */ + std::tuple<Callbacks...> callbacks_; Review Comment: consider use a single callback, and use the chain trick(better than tail recursion) to generate such a callback ########## include/tvm/ffi/extra/structural_visit.h: ########## @@ -0,0 +1,535 @@ +/* + * 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_visit.h + * \brief Structural visit implementation + */ +#ifndef TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_ +#define TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_ + +#include <tvm/ffi/any.h> +#include <tvm/ffi/c_api.h> +#include <tvm/ffi/container/variant.h> +#include <tvm/ffi/expected.h> +#include <tvm/ffi/function_details.h> +#include <tvm/ffi/optional.h> + +#include <cstddef> +#include <exception> +#include <optional> +#include <tuple> +#include <type_traits> +#include <utility> + +namespace tvm { +namespace ffi { + +/*! + * \brief Object node carrying the optional payload for an interrupted structural visit. + */ +class VisitInterruptObj : public Object { + public: + /*! \brief Payload returned with the interrupt, or FFI None for no payload. */ + Any value; + + VisitInterruptObj() = default; + /*! + * \brief Construct a VisitInterruptObj with a payload. + * \param value The payload carried by the interrupt. + */ + explicit VisitInterruptObj(Any value) : value(std::move(value)) {} + + /// \cond Doxygen_Suppress + static constexpr const int32_t _type_index = TypeIndex::kTVMFFIVisitInterrupt; + static const constexpr bool _type_final = true; + TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIVisitInterrupt, VisitInterruptObj, + Object); + /// \endcond +}; + +/*! + * \brief ObjectRef wrapper for VisitInterruptObj. + */ +class VisitInterrupt : public ObjectRef { + public: + /*! \brief Construct an interrupt with no payload. */ + VisitInterrupt() : VisitInterrupt(Any(nullptr)) {} + /*! + * \brief Construct an interrupt with a user-defined payload. + * \param value The payload carried by the interrupt. + */ + explicit VisitInterrupt(Any value) + : ObjectRef(make_object<VisitInterruptObj>(std::move(value))) {} + + /// \cond Doxygen_Suppress + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VisitInterrupt, ObjectRef, VisitInterruptObj); + /// \endcond +}; + +class StructuralVisitorObj; + +/*! + * \brief C++ vtable for structural visitor dispatch. + * + * The vtable gives \ref StructuralVisitorObj an explicit dispatch table instead + * of relying on C++ virtual methods. + */ +struct StructuralVisitorVTable { + /*! \brief Visit callback. */ + TVMFFIAny (*visit)(StructuralVisitorObj* visitor, const Object* value) noexcept = nullptr; +}; + +/*! + * \brief Object node of a structural visitor. + * + * A structural visitor is an active traversal context. It carries the dispatch + * table used to visit each object and the current def-region state used by + * structural equality/hash semantics. The visitor is ref-counted so it can + * cross FFI boundaries, but one underlying visitor object should not be shared + * by overlapping top-level traversals. + */ +class StructuralVisitorObj : public Object { + public: + /*! \brief Construct the default structural visitor. */ + StructuralVisitorObj() : StructuralVisitorObj(VTable()) {} + + /*! + * \brief Visit a value, dispatching through this visitor's vtable. + * + * \param value The object to visit. + * \return ``std::nullopt`` to continue traversal, or a \ref VisitInterrupt + * to halt the entire visit. + */ + TVM_FFI_INLINE Optional<VisitInterrupt> Visit(const ObjectRef& value) { + return VisitExpected(value).value(); + } + + /*! + * \brief Visit a value, propagating error through expected return. + * + * \param value The object to visit. + * \return Expected interrupt state. An error means traversal failed. + */ + TVM_FFI_INLINE Expected<Optional<VisitInterrupt>> VisitExpected(const ObjectRef& value) noexcept { + return details::ExpectedUnsafe::MoveFromTVMFFIAny<Optional<VisitInterrupt>>( + (*vtable_->visit)(this, value.get())); + } + + /*! + * \brief Default visit behavior: type-attr lookup, then reflection fallback. + * + * The default path first checks the structural-visit type attribute for a + * type-specific override. If none is registered, it visits all reflected + * fields that participate in structural equality/hash. + */ + TVM_FFI_INLINE Optional<VisitInterrupt> DefaultVisit(const ObjectRef& value) { + return DefaultVisitExpected(value).value(); + } + + /*! + * \brief Default visit behavior, propagating error through expected return. + * + * \param value The object to visit. + * \return Expected interrupt state. An error means traversal failed. + */ + TVM_FFI_DLL Expected<Optional<VisitInterrupt>> DefaultVisitExpected( + const ObjectRef& value) noexcept; + + /*! + * \brief Temporarily switch the def-region context while invoking \p callback. + * + * This helper scopes updates to the traversal state used by def/use-region + * aware visitors. The previous state is restored when the callback returns + * or throws. + * + * \param kind The def-region kind to set during the callback. + * \param callback A nullary callable that performs recursive visiting. + * \return The value returned by \p callback. + */ + template <typename Callback> + TVM_FFI_INLINE auto WithDefRegionKind(TVMFFIDefRegionKind kind, Callback&& callback) { + class Scope { + public: + Scope(StructuralVisitorObj* visitor, TVMFFIDefRegionKind kind) + : visitor_(visitor), old_kind_(visitor->def_region_mode_) { + visitor_->def_region_mode_ = kind; + } + ~Scope() { visitor_->def_region_mode_ = old_kind_; } + Scope(const Scope&) = delete; + Scope& operator=(const Scope&) = delete; + + private: + StructuralVisitorObj* visitor_; + TVMFFIDefRegionKind old_kind_; + }; + Scope scope(this, kind); + return std::forward<Callback>(callback)(); + } + + /// \cond Doxygen_Suppress + static constexpr const bool _type_mutable = true; + TVM_FFI_DECLARE_OBJECT_INFO("ffi.StructuralVisitor", StructuralVisitorObj, Object); + /// \endcond + + protected: + /*! + * \brief Construct a structural visitor subclass with a custom dispatch vtable. + * + * \param vtable The non-null dispatch table for this visitor. + * + * \note This constructor is for internal subclasses. The vtable and its + * ``visit`` callback must be valid for the lifetime of the visitor. + */ + explicit StructuralVisitorObj(const StructuralVisitorVTable* vtable) : vtable_(vtable) {} + + /*! + * \brief Required ABI dispatch table. \ref StructuralVisitorVTable + * It must never be null on a constructed visitor. + */ + const StructuralVisitorVTable* vtable_ = nullptr; + + /*! + * \brief Current def-region context for structural equality/hash semantics. + * + * This is shared mutable traversal state. Be careful when mutating it through + * multiple references to the same visitor object. Use \ref WithDefRegionKind + * to scope temporary changes. + */ + TVMFFIDefRegionKind def_region_mode_ = kTVMFFIDefRegionKindNone; + + private: + /*! + * \brief Return the vtable used by the default visitor implementation. + * \return Pointer to the static structural visitor vtable. + */ + static const StructuralVisitorVTable* VTable() { + static const StructuralVisitorVTable vtable{&StructuralVisitorObj::DispatchVisit}; + return &vtable; + } + + /*! + * \brief Dispatch from the vtable to the default visitor implementation. + * \param visitor The structural visitor object. + * \param value The object to visit. + * \return Interrupt state, or an error if traversal failed. + */ + static TVMFFIAny DispatchVisit(StructuralVisitorObj* visitor, const Object* value) noexcept { + ObjectRef value_ref( + details::ObjectUnsafe::ObjectPtrFromUnowned<Object>(const_cast<Object*>(value))); + return details::ExpectedUnsafe::MoveToTVMFFIAny(visitor->DefaultVisitExpected(value_ref)); + } +}; + +/*! + * \brief ObjectRef wrapper of \ref StructuralVisitorObj. + * + * \sa StructuralVisitorObj + */ +class StructuralVisitor : public ObjectRef { + public: + /*! + * \brief Construct the default structural visitor. + */ + StructuralVisitor() : ObjectRef(make_object<StructuralVisitorObj>()) {} + /*! + * \brief Construct from an existing object pointer. + * \param n The object pointer to wrap. + */ + explicit StructuralVisitor(ObjectPtr<StructuralVisitorObj> n) : ObjectRef(std::move(n)) {} + + /// \cond Doxygen_Suppress + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StructuralVisitor, ObjectRef, StructuralVisitorObj); + /// \endcond +}; + +// --------------------------------------------------------------------------- +// Structural Walk API. +// --------------------------------------------------------------------------- + +/*! + * \brief Callback order for \ref StructuralWalk. + */ +enum class WalkOrder : int32_t { + /*! \brief Invoke the callback before visiting children. */ + kPreOrder = 0, + /*! \brief Invoke the callback after visiting children. */ + kPostOrder = 1, +}; + +namespace details { + +/*! + * \brief Return true when \p result already carries a traversal-stopping state. + * \tparam T The Expected success type. + * \param result The Expected value to inspect. + * \return Whether \p result stores an Error or VisitInterrupt. + */ +template <typename T> +TVM_FFI_INLINE bool NeedEarlyReturn(const Expected<T>& result) noexcept { + int32_t type_index = result.type_index(); + return type_index == TypeIndex::kTVMFFIError || type_index == TypeIndex::kTVMFFIVisitInterrupt; +} + +/*! + * \brief Return from the current ABI visit function if \p Result stops traversal. + * + * \param Result Expression returning an Expected value whose raw storage can be + * moved to TVMFFIAny. + */ +#define TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Result) \ + if (auto result = (Result); \ + TVM_FFI_PREDICT_FALSE(::tvm::ffi::details::NeedEarlyReturn(result))) { \ + return ::tvm::ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result)); \ + } + +/*! + * \brief Move raw Expected storage into another Expected type. + * \tparam T The destination Expected success type. + * \tparam U The source Expected success type. + * \param result The source Expected value to move from. + * \return Expected<T> backed by the moved raw storage. + * + * \note This assumes the stored value is ABI-compatible with T, or is Error. + */ +template <typename T, typename U> +TVM_FFI_INLINE Expected<T> UnsafeMoveExpected(Expected<U>&& result) { + return ExpectedUnsafe::MoveFromTVMFFIAny<T>(ExpectedUnsafe::MoveToTVMFFIAny(std::move(result))); +} + +/*! + * \brief Read an Expected success value as a compatible raw storage type. + * \tparam T The type to read from the underlying Any storage. + * \tparam U The Expected success type. + * \param result The Expected value to read from. + * \return The stored value decoded as T. + * + * \note This assumes \p result stores T-compatible Any storage, or Error. + */ +template <typename T, typename U> +TVM_FFI_INLINE T UnsafeExpectedValueAs(const Expected<U>& result) { + const Any& data = ExpectedUnsafe::GetData(result); + if (TVM_FFI_PREDICT_TRUE(data.type_index() != TypeIndex::kTVMFFIError)) { + return AnyUnsafe::CopyFromAnyViewAfterCheck<T>(data); + } + throw AnyUnsafe::CopyFromAnyViewAfterCheck<Error>(data); +} + +/*! + * \brief Concrete visitor implementation used by callback-dispatched ``StructuralWalk``. + * \tparam order Callback placement relative to child traversal. + * \tparam Callbacks Callback types to invoke on matching nodes. + */ +template <WalkOrder order, typename... Callbacks> +class StructuralWalkCallbackVisitorObj : public StructuralVisitorObj { + public: + /*! + * \brief Construct a structural walk visitor. + * \param callbacks Callbacks to invoke on matching nodes. + */ + template <typename... F> + explicit StructuralWalkCallbackVisitorObj(F&&... callbacks) + : StructuralVisitorObj(VTable()), callbacks_(std::forward<F>(callbacks)...) {} + + private: + /*! + * \brief Return the vtable used by this visitor implementation. + * \return Pointer to the static structural visitor vtable. + */ + static const StructuralVisitorVTable* VTable() { + static const StructuralVisitorVTable vtable{&StructuralWalkCallbackVisitorObj::DispatchVisit}; + return &vtable; + } + + /*! + * \brief Dispatch from the erased visitor pointer to the concrete walk visitor. + * \param self The erased structural visitor object. + * \param value The object to visit. + * \return Interrupt state, or an error if traversal failed. + */ + static TVMFFIAny DispatchVisit(StructuralVisitorObj* self, const Object* value) noexcept { + ObjectRef value_ref( + details::ObjectUnsafe::ObjectPtrFromUnowned<Object>(const_cast<Object*>(value))); Review Comment: You can do GetRef<ObjectRef>(value) for such promotio ########## include/tvm/ffi/extra/structural_visit.h: ########## @@ -0,0 +1,535 @@ +/* + * 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_visit.h + * \brief Structural visit implementation + */ +#ifndef TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_ +#define TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_ + +#include <tvm/ffi/any.h> +#include <tvm/ffi/c_api.h> +#include <tvm/ffi/container/variant.h> +#include <tvm/ffi/expected.h> +#include <tvm/ffi/function_details.h> +#include <tvm/ffi/optional.h> + +#include <cstddef> +#include <exception> +#include <optional> +#include <tuple> +#include <type_traits> +#include <utility> + +namespace tvm { +namespace ffi { + +/*! + * \brief Object node carrying the optional payload for an interrupted structural visit. + */ +class VisitInterruptObj : public Object { + public: + /*! \brief Payload returned with the interrupt, or FFI None for no payload. */ + Any value; + + VisitInterruptObj() = default; + /*! + * \brief Construct a VisitInterruptObj with a payload. + * \param value The payload carried by the interrupt. + */ + explicit VisitInterruptObj(Any value) : value(std::move(value)) {} + + /// \cond Doxygen_Suppress + static constexpr const int32_t _type_index = TypeIndex::kTVMFFIVisitInterrupt; + static const constexpr bool _type_final = true; + TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIVisitInterrupt, VisitInterruptObj, + Object); + /// \endcond +}; + +/*! + * \brief ObjectRef wrapper for VisitInterruptObj. + */ +class VisitInterrupt : public ObjectRef { + public: + /*! \brief Construct an interrupt with no payload. */ + VisitInterrupt() : VisitInterrupt(Any(nullptr)) {} + /*! + * \brief Construct an interrupt with a user-defined payload. + * \param value The payload carried by the interrupt. + */ + explicit VisitInterrupt(Any value) + : ObjectRef(make_object<VisitInterruptObj>(std::move(value))) {} + + /// \cond Doxygen_Suppress + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VisitInterrupt, ObjectRef, VisitInterruptObj); + /// \endcond +}; + +class StructuralVisitorObj; + +/*! + * \brief C++ vtable for structural visitor dispatch. + * + * The vtable gives \ref StructuralVisitorObj an explicit dispatch table instead + * of relying on C++ virtual methods. + */ +struct StructuralVisitorVTable { + /*! \brief Visit callback. */ + TVMFFIAny (*visit)(StructuralVisitorObj* visitor, const Object* value) noexcept = nullptr; +}; + +/*! + * \brief Object node of a structural visitor. + * + * A structural visitor is an active traversal context. It carries the dispatch + * table used to visit each object and the current def-region state used by + * structural equality/hash semantics. The visitor is ref-counted so it can + * cross FFI boundaries, but one underlying visitor object should not be shared + * by overlapping top-level traversals. + */ +class StructuralVisitorObj : public Object { + public: + /*! \brief Construct the default structural visitor. */ + StructuralVisitorObj() : StructuralVisitorObj(VTable()) {} + + /*! + * \brief Visit a value, dispatching through this visitor's vtable. + * + * \param value The object to visit. + * \return ``std::nullopt`` to continue traversal, or a \ref VisitInterrupt + * to halt the entire visit. + */ + TVM_FFI_INLINE Optional<VisitInterrupt> Visit(const ObjectRef& value) { + return VisitExpected(value).value(); + } + + /*! + * \brief Visit a value, propagating error through expected return. + * + * \param value The object to visit. + * \return Expected interrupt state. An error means traversal failed. + */ + TVM_FFI_INLINE Expected<Optional<VisitInterrupt>> VisitExpected(const ObjectRef& value) noexcept { + return details::ExpectedUnsafe::MoveFromTVMFFIAny<Optional<VisitInterrupt>>( + (*vtable_->visit)(this, value.get())); + } + + /*! + * \brief Default visit behavior: type-attr lookup, then reflection fallback. + * + * The default path first checks the structural-visit type attribute for a + * type-specific override. If none is registered, it visits all reflected + * fields that participate in structural equality/hash. + */ + TVM_FFI_INLINE Optional<VisitInterrupt> DefaultVisit(const ObjectRef& value) { + return DefaultVisitExpected(value).value(); + } + + /*! + * \brief Default visit behavior, propagating error through expected return. + * + * \param value The object to visit. + * \return Expected interrupt state. An error means traversal failed. + */ + TVM_FFI_DLL Expected<Optional<VisitInterrupt>> DefaultVisitExpected( + const ObjectRef& value) noexcept; + + /*! + * \brief Temporarily switch the def-region context while invoking \p callback. + * + * This helper scopes updates to the traversal state used by def/use-region + * aware visitors. The previous state is restored when the callback returns + * or throws. + * + * \param kind The def-region kind to set during the callback. + * \param callback A nullary callable that performs recursive visiting. + * \return The value returned by \p callback. + */ + template <typename Callback> + TVM_FFI_INLINE auto WithDefRegionKind(TVMFFIDefRegionKind kind, Callback&& callback) { + class Scope { + public: + Scope(StructuralVisitorObj* visitor, TVMFFIDefRegionKind kind) + : visitor_(visitor), old_kind_(visitor->def_region_mode_) { + visitor_->def_region_mode_ = kind; + } + ~Scope() { visitor_->def_region_mode_ = old_kind_; } + Scope(const Scope&) = delete; + Scope& operator=(const Scope&) = delete; + + private: + StructuralVisitorObj* visitor_; + TVMFFIDefRegionKind old_kind_; + }; + Scope scope(this, kind); + return std::forward<Callback>(callback)(); + } + + /// \cond Doxygen_Suppress + static constexpr const bool _type_mutable = true; + TVM_FFI_DECLARE_OBJECT_INFO("ffi.StructuralVisitor", StructuralVisitorObj, Object); + /// \endcond + + protected: + /*! + * \brief Construct a structural visitor subclass with a custom dispatch vtable. + * + * \param vtable The non-null dispatch table for this visitor. + * + * \note This constructor is for internal subclasses. The vtable and its + * ``visit`` callback must be valid for the lifetime of the visitor. + */ + explicit StructuralVisitorObj(const StructuralVisitorVTable* vtable) : vtable_(vtable) {} + + /*! + * \brief Required ABI dispatch table. \ref StructuralVisitorVTable + * It must never be null on a constructed visitor. + */ + const StructuralVisitorVTable* vtable_ = nullptr; + + /*! + * \brief Current def-region context for structural equality/hash semantics. + * + * This is shared mutable traversal state. Be careful when mutating it through + * multiple references to the same visitor object. Use \ref WithDefRegionKind + * to scope temporary changes. + */ + TVMFFIDefRegionKind def_region_mode_ = kTVMFFIDefRegionKindNone; + + private: + /*! + * \brief Return the vtable used by the default visitor implementation. + * \return Pointer to the static structural visitor vtable. + */ + static const StructuralVisitorVTable* VTable() { + static const StructuralVisitorVTable vtable{&StructuralVisitorObj::DispatchVisit}; + return &vtable; + } + + /*! + * \brief Dispatch from the vtable to the default visitor implementation. + * \param visitor The structural visitor object. + * \param value The object to visit. + * \return Interrupt state, or an error if traversal failed. + */ + static TVMFFIAny DispatchVisit(StructuralVisitorObj* visitor, const Object* value) noexcept { + ObjectRef value_ref( + details::ObjectUnsafe::ObjectPtrFromUnowned<Object>(const_cast<Object*>(value))); + return details::ExpectedUnsafe::MoveToTVMFFIAny(visitor->DefaultVisitExpected(value_ref)); + } +}; + +/*! + * \brief ObjectRef wrapper of \ref StructuralVisitorObj. + * + * \sa StructuralVisitorObj + */ +class StructuralVisitor : public ObjectRef { + public: + /*! + * \brief Construct the default structural visitor. + */ + StructuralVisitor() : ObjectRef(make_object<StructuralVisitorObj>()) {} + /*! + * \brief Construct from an existing object pointer. + * \param n The object pointer to wrap. + */ + explicit StructuralVisitor(ObjectPtr<StructuralVisitorObj> n) : ObjectRef(std::move(n)) {} + + /// \cond Doxygen_Suppress + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StructuralVisitor, ObjectRef, StructuralVisitorObj); + /// \endcond +}; + +// --------------------------------------------------------------------------- +// Structural Walk API. +// --------------------------------------------------------------------------- + +/*! + * \brief Callback order for \ref StructuralWalk. + */ +enum class WalkOrder : int32_t { + /*! \brief Invoke the callback before visiting children. */ + kPreOrder = 0, + /*! \brief Invoke the callback after visiting children. */ + kPostOrder = 1, +}; + +namespace details { + +/*! + * \brief Return true when \p result already carries a traversal-stopping state. + * \tparam T The Expected success type. + * \param result The Expected value to inspect. + * \return Whether \p result stores an Error or VisitInterrupt. + */ +template <typename T> +TVM_FFI_INLINE bool NeedEarlyReturn(const Expected<T>& result) noexcept { + int32_t type_index = result.type_index(); + return type_index == TypeIndex::kTVMFFIError || type_index == TypeIndex::kTVMFFIVisitInterrupt; +} + +/*! + * \brief Return from the current ABI visit function if \p Result stops traversal. + * + * \param Result Expression returning an Expected value whose raw storage can be + * moved to TVMFFIAny. + */ +#define TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Result) \ + if (auto result = (Result); \ + TVM_FFI_PREDICT_FALSE(::tvm::ffi::details::NeedEarlyReturn(result))) { \ + return ::tvm::ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result)); \ + } + +/*! + * \brief Move raw Expected storage into another Expected type. + * \tparam T The destination Expected success type. + * \tparam U The source Expected success type. + * \param result The source Expected value to move from. + * \return Expected<T> backed by the moved raw storage. + * + * \note This assumes the stored value is ABI-compatible with T, or is Error. + */ +template <typename T, typename U> +TVM_FFI_INLINE Expected<T> UnsafeMoveExpected(Expected<U>&& result) { + return ExpectedUnsafe::MoveFromTVMFFIAny<T>(ExpectedUnsafe::MoveToTVMFFIAny(std::move(result))); +} + +/*! + * \brief Read an Expected success value as a compatible raw storage type. + * \tparam T The type to read from the underlying Any storage. + * \tparam U The Expected success type. + * \param result The Expected value to read from. + * \return The stored value decoded as T. + * + * \note This assumes \p result stores T-compatible Any storage, or Error. + */ +template <typename T, typename U> +TVM_FFI_INLINE T UnsafeExpectedValueAs(const Expected<U>& result) { + const Any& data = ExpectedUnsafe::GetData(result); + if (TVM_FFI_PREDICT_TRUE(data.type_index() != TypeIndex::kTVMFFIError)) { + return AnyUnsafe::CopyFromAnyViewAfterCheck<T>(data); + } + throw AnyUnsafe::CopyFromAnyViewAfterCheck<Error>(data); +} + +/*! + * \brief Concrete visitor implementation used by callback-dispatched ``StructuralWalk``. + * \tparam order Callback placement relative to child traversal. + * \tparam Callbacks Callback types to invoke on matching nodes. + */ +template <WalkOrder order, typename... Callbacks> +class StructuralWalkCallbackVisitorObj : public StructuralVisitorObj { + public: + /*! + * \brief Construct a structural walk visitor. + * \param callbacks Callbacks to invoke on matching nodes. + */ + template <typename... F> + explicit StructuralWalkCallbackVisitorObj(F&&... callbacks) + : StructuralVisitorObj(VTable()), callbacks_(std::forward<F>(callbacks)...) {} + + private: + /*! + * \brief Return the vtable used by this visitor implementation. + * \return Pointer to the static structural visitor vtable. + */ + static const StructuralVisitorVTable* VTable() { + static const StructuralVisitorVTable vtable{&StructuralWalkCallbackVisitorObj::DispatchVisit}; + return &vtable; + } + + /*! + * \brief Dispatch from the erased visitor pointer to the concrete walk visitor. + * \param self The erased structural visitor object. + * \param value The object to visit. + * \return Interrupt state, or an error if traversal failed. + */ + static TVMFFIAny DispatchVisit(StructuralVisitorObj* self, const Object* value) noexcept { + ObjectRef value_ref( + details::ObjectUnsafe::ObjectPtrFromUnowned<Object>(const_cast<Object*>(value))); + return details::ExpectedUnsafe::MoveToTVMFFIAny( + static_cast<StructuralWalkCallbackVisitorObj*>(self)->VisitImpl(value_ref)); + } + + /*! + * \brief Visit one object according to the configured walk order. + * \param value The object to visit. + * \return Interrupt state, or an error if traversal failed. + */ + Expected<Optional<VisitInterrupt>> VisitImpl(const ObjectRef& value) noexcept { + try { + if constexpr (order == WalkOrder::kPreOrder) { + auto result = InvokeMatchingCallback(value); + if (TVM_FFI_PREDICT_FALSE(NeedEarlyReturn(result))) { + return UnsafeMoveExpected<Optional<VisitInterrupt>>(std::move(result)); + } + int32_t type_index = result.type_index(); + TVM_FFI_UNSAFE_ASSUME(type_index == TypeIndex::kTVMFFIBool); + if (TVM_FFI_PREDICT_FALSE(!UnsafeExpectedValueAs<bool>(result))) { + return Optional<VisitInterrupt>(std::nullopt); + } + } + + Expected<Optional<VisitInterrupt>> interrupt = DefaultVisitExpected(value); + if (TVM_FFI_PREDICT_FALSE(NeedEarlyReturn(interrupt))) { + return interrupt; + } + + if constexpr (order == WalkOrder::kPostOrder) { + auto result = InvokeMatchingCallback(value); + if (TVM_FFI_PREDICT_FALSE(NeedEarlyReturn(result))) { + return UnsafeMoveExpected<Optional<VisitInterrupt>>(std::move(result)); + } + } + + return Optional<VisitInterrupt>(std::nullopt); + } catch (const Error& err) { + return Unexpected(err); + } catch (const std::exception& err) { + return Unexpected(Error("InternalError", err.what(), "")); + } catch (...) { + return Unexpected(Error("InternalError", "Unknown structural walk error", "")); + } + } + + /*! + * \brief Invoke the first callback whose argument type matches ``value``. + * \param value The object to match against the callback arguments. + * \return Callback action, or an error if the callback failed. + */ + Expected<Variant<VisitInterrupt, bool>> InvokeMatchingCallback(const ObjectRef& value) { + if (TVM_FFI_PREDICT_FALSE(!value.defined())) { + return Variant<VisitInterrupt, bool>(true); + } + return InvokeMatchingCallback<0>(value); + } + + /*! + * \brief Invoke the first matching callback from index ``I`` onward. + * \param value The object to match. + * \return Callback action, or continue if no callback matched. + */ + template <std::size_t I> + Expected<Variant<VisitInterrupt, bool>> InvokeMatchingCallback(const ObjectRef& value) { + if constexpr (I == sizeof...(Callbacks)) { + return Variant<VisitInterrupt, bool>(true); + } else { + using Callback = std::tuple_element_t<I, std::tuple<Callbacks...>>; + using FuncInfo = FunctionInfo<Callback>; + using Arg = std::tuple_element_t<0, typename FuncInfo::ArgType>; + using RefType = std::remove_cv_t<std::remove_reference_t<Arg>>; + + if (auto ref = value.template as<RefType>()) { + return std::get<I>(callbacks_)(*ref); + } + return InvokeMatchingCallback<I + 1>(value); + } + } + + /*! \brief Concrete callback objects. */ + std::tuple<Callbacks...> callbacks_; +}; + +} // namespace details + +/*! + * \brief Walk an object graph structurally and invoke typed callbacks on selected node types. + * + * The callbacks are invoked only for nodes matching the first argument type of + * one of the callbacks. Callback arguments must be ``const ObjectRef&`` or + * ``const`` references to subclasses of ``ObjectRef``. Callbacks are tested in + * order, and the first match is used. + * + * Each callback should return ``Expected<Variant<VisitInterrupt, bool>>``: + * - ``VisitInterrupt`` halts traversal, analogous to interrupt in MLIR. + * - ``true`` continues traversal, analogous to advance in MLIR. + * - ``false`` skips children traversal, analogous to skip in MLIR. + * - ``Error`` indicates traversal failure. + * See \ref WalkOrder for more details. + * + * Example: + * + * \code + * int num_adds = 0; + * + * Expected<Optional<VisitInterrupt>> result = StructuralWalkExpected<WalkOrder::kPreOrder>( + * root, + * [&](const Add& add) -> Expected<Variant<VisitInterrupt, bool>> { + * ++num_adds; + * return Variant<VisitInterrupt, bool>(true); + * }, + * [&](const Mul& mul) -> Expected<Variant<VisitInterrupt, bool>> { + * return Variant<VisitInterrupt, bool>(false); + * }); + * \endcode + * + * \tparam order Whether to invoke the callback before or after visiting children. + * \tparam Callbacks Callback types. + * \param root The root object to visit. + * \param callbacks Callbacks invoked for matching nodes. + * \return ``std::nullopt`` if traversal completed, or the interrupt returned by + * a callback. + * + * \note Return type of each callback should be + * Expected<Variant<VisitInterrupt, bool>>. + */ +template <WalkOrder order, typename... Callbacks> +Expected<Optional<VisitInterrupt>> StructuralWalkExpected(const ObjectRef& root, + Callbacks&&... callbacks) noexcept { + try { + static_assert(sizeof...(Callbacks) != 0, "StructuralWalk requires at least one callback"); + + using Visitor = details::StructuralWalkCallbackVisitorObj<order, std::decay_t<Callbacks>...>; + StructuralVisitor visitor(make_object<Visitor>(std::forward<Callbacks>(callbacks)...)); + return visitor->VisitExpected(root); Review Comment: if we are VisitExpected, it shiould be noexcept so no need to try catch around VisitExpected ########## include/tvm/ffi/expected.h: ########## @@ -186,9 +195,106 @@ class Expected { } private: + Expected() = default; + + friend struct details::ExpectedUnsafe; + Any data_; // Invariant: holds a T (type_index != kTVMFFIError) or an Error. }; +namespace details { + +/*! + * \brief Unsafe raw-storage helpers for Expected. + * + * These helpers bypass normal value checking and are intended for ABI boundaries + * that already know the underlying Any storage holds either a valid T or Error. + */ +struct ExpectedUnsafe { + /*! + * \brief Copy a raw TVMFFIAny into Expected storage. + * \tparam T The Expected success type. + * \param raw The raw FFI value to copy. + * \return Expected backed by a copied Any storage. + */ + template <typename T> + TVM_FFI_INLINE static Expected<T> CopyFromTVMFFIAny(TVMFFIAny raw) { + Expected<T> result; + result.data_ = Any(AnyView::CopyFromTVMFFIAny(raw)); + return result; + } + + /*! + * \brief Move a raw TVMFFIAny into Expected storage. + * \tparam T The Expected success type. + * \param raw The raw FFI value to move. + * \return Expected backed by moved Any storage. + */ + template <typename T> + TVM_FFI_INLINE static Expected<T> MoveFromTVMFFIAny(TVMFFIAny raw) { + Expected<T> result; + result.data_ = AnyUnsafe::MoveTVMFFIAnyToAny(&raw); + return result; + } + + /*! + * \brief Copy Expected storage to a raw TVMFFIAny. + * \tparam T The Expected success type. + * \param result The Expected value to copy from. + * \return Raw FFI value containing a copy of the underlying Any storage. + */ + template <typename T> + TVM_FFI_INLINE static TVMFFIAny CopyToTVMFFIAny(const Expected<T>& result) { Review Comment: this is not used, don't expose ########## include/tvm/ffi/extra/structural_visit.h: ########## @@ -0,0 +1,535 @@ +/* + * 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_visit.h + * \brief Structural visit implementation + */ +#ifndef TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_ +#define TVM_FFI_EXTRA_STRUCTURAL_VISIT_H_ + +#include <tvm/ffi/any.h> +#include <tvm/ffi/c_api.h> +#include <tvm/ffi/container/variant.h> +#include <tvm/ffi/expected.h> +#include <tvm/ffi/function_details.h> +#include <tvm/ffi/optional.h> + +#include <cstddef> +#include <exception> +#include <optional> +#include <tuple> +#include <type_traits> +#include <utility> + +namespace tvm { +namespace ffi { + +/*! + * \brief Object node carrying the optional payload for an interrupted structural visit. + */ +class VisitInterruptObj : public Object { + public: + /*! \brief Payload returned with the interrupt, or FFI None for no payload. */ + Any value; + + VisitInterruptObj() = default; + /*! + * \brief Construct a VisitInterruptObj with a payload. + * \param value The payload carried by the interrupt. + */ + explicit VisitInterruptObj(Any value) : value(std::move(value)) {} + + /// \cond Doxygen_Suppress + static constexpr const int32_t _type_index = TypeIndex::kTVMFFIVisitInterrupt; + static const constexpr bool _type_final = true; + TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFIVisitInterrupt, VisitInterruptObj, + Object); + /// \endcond +}; + +/*! + * \brief ObjectRef wrapper for VisitInterruptObj. + */ +class VisitInterrupt : public ObjectRef { + public: + /*! \brief Construct an interrupt with no payload. */ + VisitInterrupt() : VisitInterrupt(Any(nullptr)) {} + /*! + * \brief Construct an interrupt with a user-defined payload. + * \param value The payload carried by the interrupt. + */ + explicit VisitInterrupt(Any value) + : ObjectRef(make_object<VisitInterruptObj>(std::move(value))) {} + + /// \cond Doxygen_Suppress + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VisitInterrupt, ObjectRef, VisitInterruptObj); + /// \endcond +}; + +class StructuralVisitorObj; + +/*! + * \brief C++ vtable for structural visitor dispatch. + * + * The vtable gives \ref StructuralVisitorObj an explicit dispatch table instead + * of relying on C++ virtual methods. + */ +struct StructuralVisitorVTable { + /*! \brief Visit callback. */ + TVMFFIAny (*visit)(StructuralVisitorObj* visitor, const Object* value) noexcept = nullptr; +}; + +/*! + * \brief Object node of a structural visitor. + * + * A structural visitor is an active traversal context. It carries the dispatch + * table used to visit each object and the current def-region state used by + * structural equality/hash semantics. The visitor is ref-counted so it can + * cross FFI boundaries, but one underlying visitor object should not be shared + * by overlapping top-level traversals. + */ +class StructuralVisitorObj : public Object { + public: + /*! \brief Construct the default structural visitor. */ + StructuralVisitorObj() : StructuralVisitorObj(VTable()) {} + + /*! + * \brief Visit a value, dispatching through this visitor's vtable. + * + * \param value The object to visit. + * \return ``std::nullopt`` to continue traversal, or a \ref VisitInterrupt + * to halt the entire visit. + */ + TVM_FFI_INLINE Optional<VisitInterrupt> Visit(const ObjectRef& value) { + return VisitExpected(value).value(); + } + + /*! + * \brief Visit a value, propagating error through expected return. + * + * \param value The object to visit. + * \return Expected interrupt state. An error means traversal failed. + */ + TVM_FFI_INLINE Expected<Optional<VisitInterrupt>> VisitExpected(const ObjectRef& value) noexcept { + return details::ExpectedUnsafe::MoveFromTVMFFIAny<Optional<VisitInterrupt>>( + (*vtable_->visit)(this, value.get())); + } + + /*! + * \brief Default visit behavior: type-attr lookup, then reflection fallback. + * + * The default path first checks the structural-visit type attribute for a + * type-specific override. If none is registered, it visits all reflected + * fields that participate in structural equality/hash. + */ + TVM_FFI_INLINE Optional<VisitInterrupt> DefaultVisit(const ObjectRef& value) { + return DefaultVisitExpected(value).value(); + } + + /*! + * \brief Default visit behavior, propagating error through expected return. + * + * \param value The object to visit. + * \return Expected interrupt state. An error means traversal failed. + */ + TVM_FFI_DLL Expected<Optional<VisitInterrupt>> DefaultVisitExpected( + const ObjectRef& value) noexcept; + + /*! + * \brief Temporarily switch the def-region context while invoking \p callback. + * + * This helper scopes updates to the traversal state used by def/use-region + * aware visitors. The previous state is restored when the callback returns + * or throws. + * + * \param kind The def-region kind to set during the callback. + * \param callback A nullary callable that performs recursive visiting. + * \return The value returned by \p callback. + */ + template <typename Callback> + TVM_FFI_INLINE auto WithDefRegionKind(TVMFFIDefRegionKind kind, Callback&& callback) { + class Scope { + public: + Scope(StructuralVisitorObj* visitor, TVMFFIDefRegionKind kind) + : visitor_(visitor), old_kind_(visitor->def_region_mode_) { + visitor_->def_region_mode_ = kind; + } + ~Scope() { visitor_->def_region_mode_ = old_kind_; } + Scope(const Scope&) = delete; + Scope& operator=(const Scope&) = delete; + + private: + StructuralVisitorObj* visitor_; + TVMFFIDefRegionKind old_kind_; + }; + Scope scope(this, kind); + return std::forward<Callback>(callback)(); + } + + /// \cond Doxygen_Suppress + static constexpr const bool _type_mutable = true; + TVM_FFI_DECLARE_OBJECT_INFO("ffi.StructuralVisitor", StructuralVisitorObj, Object); + /// \endcond + + protected: + /*! + * \brief Construct a structural visitor subclass with a custom dispatch vtable. + * + * \param vtable The non-null dispatch table for this visitor. + * + * \note This constructor is for internal subclasses. The vtable and its + * ``visit`` callback must be valid for the lifetime of the visitor. + */ + explicit StructuralVisitorObj(const StructuralVisitorVTable* vtable) : vtable_(vtable) {} + + /*! + * \brief Required ABI dispatch table. \ref StructuralVisitorVTable + * It must never be null on a constructed visitor. + */ + const StructuralVisitorVTable* vtable_ = nullptr; + + /*! + * \brief Current def-region context for structural equality/hash semantics. + * + * This is shared mutable traversal state. Be careful when mutating it through + * multiple references to the same visitor object. Use \ref WithDefRegionKind + * to scope temporary changes. + */ + TVMFFIDefRegionKind def_region_mode_ = kTVMFFIDefRegionKindNone; + + private: + /*! + * \brief Return the vtable used by the default visitor implementation. + * \return Pointer to the static structural visitor vtable. + */ + static const StructuralVisitorVTable* VTable() { + static const StructuralVisitorVTable vtable{&StructuralVisitorObj::DispatchVisit}; + return &vtable; + } + + /*! + * \brief Dispatch from the vtable to the default visitor implementation. + * \param visitor The structural visitor object. + * \param value The object to visit. + * \return Interrupt state, or an error if traversal failed. + */ + static TVMFFIAny DispatchVisit(StructuralVisitorObj* visitor, const Object* value) noexcept { + ObjectRef value_ref( + details::ObjectUnsafe::ObjectPtrFromUnowned<Object>(const_cast<Object*>(value))); + return details::ExpectedUnsafe::MoveToTVMFFIAny(visitor->DefaultVisitExpected(value_ref)); + } +}; + +/*! + * \brief ObjectRef wrapper of \ref StructuralVisitorObj. + * + * \sa StructuralVisitorObj + */ +class StructuralVisitor : public ObjectRef { + public: + /*! + * \brief Construct the default structural visitor. + */ + StructuralVisitor() : ObjectRef(make_object<StructuralVisitorObj>()) {} + /*! + * \brief Construct from an existing object pointer. + * \param n The object pointer to wrap. + */ + explicit StructuralVisitor(ObjectPtr<StructuralVisitorObj> n) : ObjectRef(std::move(n)) {} + + /// \cond Doxygen_Suppress + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StructuralVisitor, ObjectRef, StructuralVisitorObj); + /// \endcond +}; + +// --------------------------------------------------------------------------- +// Structural Walk API. +// --------------------------------------------------------------------------- + +/*! + * \brief Callback order for \ref StructuralWalk. + */ +enum class WalkOrder : int32_t { + /*! \brief Invoke the callback before visiting children. */ + kPreOrder = 0, + /*! \brief Invoke the callback after visiting children. */ + kPostOrder = 1, +}; + +namespace details { + +/*! + * \brief Return true when \p result already carries a traversal-stopping state. + * \tparam T The Expected success type. + * \param result The Expected value to inspect. + * \return Whether \p result stores an Error or VisitInterrupt. + */ +template <typename T> +TVM_FFI_INLINE bool NeedEarlyReturn(const Expected<T>& result) noexcept { + int32_t type_index = result.type_index(); + return type_index == TypeIndex::kTVMFFIError || type_index == TypeIndex::kTVMFFIVisitInterrupt; +} + +/*! + * \brief Return from the current ABI visit function if \p Result stops traversal. + * + * \param Result Expression returning an Expected value whose raw storage can be + * moved to TVMFFIAny. + */ +#define TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Result) \ + if (auto result = (Result); \ + TVM_FFI_PREDICT_FALSE(::tvm::ffi::details::NeedEarlyReturn(result))) { \ + return ::tvm::ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(std::move(result)); \ + } + +/*! + * \brief Move raw Expected storage into another Expected type. + * \tparam T The destination Expected success type. + * \tparam U The source Expected success type. + * \param result The source Expected value to move from. + * \return Expected<T> backed by the moved raw storage. + * + * \note This assumes the stored value is ABI-compatible with T, or is Error. + */ +template <typename T, typename U> +TVM_FFI_INLINE Expected<T> UnsafeMoveExpected(Expected<U>&& result) { + return ExpectedUnsafe::MoveFromTVMFFIAny<T>(ExpectedUnsafe::MoveToTVMFFIAny(std::move(result))); +} + +/*! + * \brief Read an Expected success value as a compatible raw storage type. + * \tparam T The type to read from the underlying Any storage. + * \tparam U The Expected success type. + * \param result The Expected value to read from. + * \return The stored value decoded as T. + * + * \note This assumes \p result stores T-compatible Any storage, or Error. + */ +template <typename T, typename U> +TVM_FFI_INLINE T UnsafeExpectedValueAs(const Expected<U>& result) { + const Any& data = ExpectedUnsafe::GetData(result); + if (TVM_FFI_PREDICT_TRUE(data.type_index() != TypeIndex::kTVMFFIError)) { + return AnyUnsafe::CopyFromAnyViewAfterCheck<T>(data); + } + throw AnyUnsafe::CopyFromAnyViewAfterCheck<Error>(data); +} + +/*! + * \brief Concrete visitor implementation used by callback-dispatched ``StructuralWalk``. + * \tparam order Callback placement relative to child traversal. + * \tparam Callbacks Callback types to invoke on matching nodes. + */ +template <WalkOrder order, typename... Callbacks> +class StructuralWalkCallbackVisitorObj : public StructuralVisitorObj { + public: + /*! + * \brief Construct a structural walk visitor. + * \param callbacks Callbacks to invoke on matching nodes. + */ + template <typename... F> + explicit StructuralWalkCallbackVisitorObj(F&&... callbacks) + : StructuralVisitorObj(VTable()), callbacks_(std::forward<F>(callbacks)...) {} + + private: + /*! + * \brief Return the vtable used by this visitor implementation. + * \return Pointer to the static structural visitor vtable. + */ + static const StructuralVisitorVTable* VTable() { + static const StructuralVisitorVTable vtable{&StructuralWalkCallbackVisitorObj::DispatchVisit}; + return &vtable; + } + + /*! + * \brief Dispatch from the erased visitor pointer to the concrete walk visitor. + * \param self The erased structural visitor object. + * \param value The object to visit. + * \return Interrupt state, or an error if traversal failed. + */ + static TVMFFIAny DispatchVisit(StructuralVisitorObj* self, const Object* value) noexcept { + ObjectRef value_ref( + details::ObjectUnsafe::ObjectPtrFromUnowned<Object>(const_cast<Object*>(value))); + return details::ExpectedUnsafe::MoveToTVMFFIAny( + static_cast<StructuralWalkCallbackVisitorObj*>(self)->VisitImpl(value_ref)); + } + + /*! + * \brief Visit one object according to the configured walk order. + * \param value The object to visit. + * \return Interrupt state, or an error if traversal failed. + */ + Expected<Optional<VisitInterrupt>> VisitImpl(const ObjectRef& value) noexcept { Review Comment: VisitImpl shuold be private and take const Object* -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
