tlopex commented on code in PR #649: URL: https://github.com/apache/tvm-ffi/pull/649#discussion_r3745084027
########## src/ffi/extra/structural_mutate.cc: ########## @@ -0,0 +1,362 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +/*! + * \file src/ffi/extra/structural_mutate.cc + * \brief Structural mutator and structural map registration. + */ +#include <tvm/ffi/container/array.h> +#include <tvm/ffi/container/dict.h> +#include <tvm/ffi/container/list.h> +#include <tvm/ffi/container/map.h> +#include <tvm/ffi/extra/structural_mutate.h> +#include <tvm/ffi/function.h> +#include <tvm/ffi/reflection/accessor.h> +#include <tvm/ffi/reflection/registry.h> + +#include <utility> + +namespace tvm { +namespace ffi { + +namespace details { + +/*! + * \brief Runtime structural map for callback arrays. + * + * \param root The root value to map. + * \param callbacks Runtime callback entries of ``(type_index, ffi::Function)`` invoked as + * ``callback(value)``. + * \param callbacks_with_def_region_kind Runtime callback entries of ``(type_index, ffi::Function)`` + * invoked as ``callback(value, def_region_kind)``. + * \param order Integer value of \ref WalkOrder. + * \return The mapped owning value, or an Error. + */ +Expected<Any> StructuralMapExpected( + AnyView root, const Array<Tuple<int32_t, Function>>& callbacks, + const Array<Tuple<int32_t, Function>>& callbacks_with_def_region_kind, int order) noexcept { + auto dispatch = [callbacks, callbacks_with_def_region_kind]( + AnyView x, TVMFFIDefRegionKind kind) -> Expected<Any> { + for (const auto& entry : callbacks) { + int32_t type_index = entry.template get<0>(); + if (!RuntimeTypeIndexMatch(x.type_index(), type_index)) { + continue; + } + Function fn = entry.template get<1>(); + return fn.CallExpected<Any>(x); + } + for (const auto& entry : callbacks_with_def_region_kind) { + int32_t type_index = entry.template get<0>(); + if (!RuntimeTypeIndexMatch(x.type_index(), type_index)) { + continue; + } + Function fn = entry.template get<1>(); + return fn.CallExpected<Any>(x, kind); + } + return Any(x); + }; + + if (order == static_cast<int>(WalkOrder::kPreOrder)) { + using Mutator = StructuralMapMutatorObj<WalkOrder::kPreOrder, decltype(dispatch)>; + StructuralMutator mutator(make_object<Mutator>(std::move(dispatch))); + return mutator->MaybeInplaceMutateIfUniqueExpected(root); + } else { + using Mutator = StructuralMapMutatorObj<WalkOrder::kPostOrder, decltype(dispatch)>; + StructuralMutator mutator(make_object<Mutator>(std::move(dispatch))); + return mutator->MaybeInplaceMutateIfUniqueExpected(root); + } +} + +// --------------------------------------------------------------------------- +// Built-in container structural mutation. +// --------------------------------------------------------------------------- + +/*! + * \brief Structurally mutate the elements of a sequence container. + * + * \tparam SeqObj The underlying sequence object type. + * \param mutator The active structural mutator. + * \param value The borrowed sequence container. + * \param source The sequence object stored in \p value. + * \return The mutated sequence, or an Error. + */ +template <typename SeqObj> +Expected<Any> MutateSeqContainerExpected(StructuralMutatorObj* mutator, AnyView value, + const SeqObj* source) noexcept { + try { + int64_t size = static_cast<int64_t>(source->size()); + ObjectPtr<SeqObj> output = nullptr; + + for (int64_t i = 0; i < size; ++i) { + const Any& item = source->at(i); + Expected<Any> mapped_item = mutator->MutateExpected(item); + if (TVM_FFI_PREDICT_FALSE(mapped_item.is_err())) { + return Unexpected(std::move(mapped_item).error()); + } + const Any& mapped_value = details::ExpectedUnsafe::GetData(mapped_item); + + if (output == nullptr) { + if (item.same_as(mapped_value)) { + continue; + } + output = SeqObj::CreateRepeated(size, Any()); + for (int64_t j = 0; j < i; ++j) { + output->SetItem(j, source->at(j)); + } + } + output->SetItem(i, mapped_value); + } + + if (output == nullptr) { + return Any(value); + } + return Any(ObjectRef(std::move(output))); + } catch (const Error& err) { + return Unexpected(err); + } +} + +/*! + * \brief Structurally mutate the elements of a sequence container in place when safe. + * + * \tparam SeqObj The underlying sequence object type. + * \param mutator The active structural mutator. + * \param value The borrowed sequence container, which must be safe to mutate in place. + * \param target The sequence object stored in \p value. + * \return The mutated sequence, or an Error. + */ +template <typename SeqObj> +Expected<Any> MaybeInplaceMutateSeqContainerExpected(StructuralMutatorObj* mutator, AnyView value, + SeqObj* target) noexcept { + try { + for (int64_t i = 0; i < static_cast<int64_t>(target->size()); ++i) { + const Any& item = target->at(i); + Expected<Any> mapped_item = mutator->MaybeInplaceMutateIfUniqueExpected(item); + if (TVM_FFI_PREDICT_FALSE(mapped_item.is_err())) { + return Unexpected(std::move(mapped_item).error()); + } + const Any& mapped_value = details::ExpectedUnsafe::GetData(mapped_item); + + if (!item.same_as(mapped_value)) { + target->SetItem(i, mapped_value); + } + } + return Any(value); + } catch (const Error& err) { + return Unexpected(err); + } +} + +/*! + * \brief Structurally mutate the values of a map container. + * + * \tparam MapObjType The underlying map object type. + * \param mutator The active structural mutator. + * \param value The borrowed map container. + * \param source The map object stored in \p value. + * \return The mutated map, or an Error. + */ +template <typename MapObjType> +Expected<Any> MutateMapValuesExpected(StructuralMutatorObj* mutator, AnyView value, + const MapObjType* source) noexcept { + try { + ObjectPtr<Object> output = nullptr; + MapBaseObj::iterator output_it; + size_t index = 0; + + for (auto source_it = source->begin(); source_it != source->end(); ++source_it, ++index) { + const Any& old_value = source_it->second; + Expected<Any> mapped_value = mutator->MutateExpected(old_value); Review Comment: If source is a Dict, MutateExpected may call a callback that mutates the same dict through an alias. This can trigger a rehash and invalidate source_it and old_value, and it may cause a SEGV when the loop continues? -- 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]
