|
| 1 | +/* |
| 2 | + * Licensed to the Apache Software Foundation (ASF) under one |
| 3 | + * or more contributor license agreements. See the NOTICE file |
| 4 | + * distributed with this work for additional information |
| 5 | + * regarding copyright ownership. The ASF licenses this file |
| 6 | + * to you under the Apache License, Version 2.0 (the |
| 7 | + * "License"); you may not use this file except in compliance |
| 8 | + * with the License. You may obtain a copy of the License at |
| 9 | + * |
| 10 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | + * |
| 12 | + * Unless required by applicable law or agreed to in writing, |
| 13 | + * software distributed under the License is distributed on an |
| 14 | + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 15 | + * KIND, either express or implied. See the License for the |
| 16 | + * specific language governing permissions and limitations |
| 17 | + * under the License. |
| 18 | + */ |
| 19 | +/*! |
| 20 | + * \file tvm/ffi/extra/type_checker.h |
| 21 | + * \brief Runtime type schema and type checker/converter. |
| 22 | + * |
| 23 | + * TypeSchema is a runtime data class describing an FFI type via a recursive |
| 24 | + * {origin, args} structure (the C++ counterpart of the Python TypeSchema). |
| 25 | + * |
| 26 | + * TypeChecker is a pre-compiled converter built from a TypeSchema. It |
| 27 | + * provides efficient runtime type checking and conversion — the dynamic |
| 28 | + * equivalent of the compile-time `AnyView::cast<T>()` / `TypeTraits<T>`. |
| 29 | + */ |
| 30 | +#ifndef TVM_FFI_EXTRA_TYPE_CHECKER_H_ |
| 31 | +#define TVM_FFI_EXTRA_TYPE_CHECKER_H_ |
| 32 | + |
| 33 | +#include <tvm/ffi/any.h> |
| 34 | +#include <tvm/ffi/extra/base.h> |
| 35 | + |
| 36 | +#include <optional> |
| 37 | +#include <string> |
| 38 | +#include <vector> |
| 39 | + |
| 40 | +namespace tvm { |
| 41 | +namespace ffi { |
| 42 | + |
| 43 | +/*! |
| 44 | + * \brief Runtime type schema describing an FFI type. |
| 45 | + * |
| 46 | + * Mirrors the Python `TypeSchema` class. The schema is expressed as a |
| 47 | + * recursive structure with an `origin` type name and optional `args` |
| 48 | + * (type arguments for generic/container types). |
| 49 | + * |
| 50 | + * Examples of JSON representations parsed by this class: |
| 51 | + * - `{"type":"int"}` |
| 52 | + * - `{"type":"Optional","args":[{"type":"ffi.String"}]}` |
| 53 | + * - `{"type":"ffi.Array","args":[{"type":"int"}]}` |
| 54 | + * - `{"type":"ffi.Map","args":[{"type":"ffi.String"},{"type":"int"}]}` |
| 55 | + */ |
| 56 | +struct TypeSchema { |
| 57 | + /*! \brief The origin type name (e.g. "int", "ffi.Array", "Optional"). */ |
| 58 | + std::string origin; |
| 59 | + /*! \brief Recursive type arguments (e.g. element type for Array). */ |
| 60 | + std::vector<TypeSchema> args; |
| 61 | + |
| 62 | + /*! |
| 63 | + * \brief Construct a TypeSchema from a parsed JSON value. |
| 64 | + * \param obj A json::Object (Map<Any,Any>) with key "type" and optional "args". |
| 65 | + * \return The constructed TypeSchema. |
| 66 | + */ |
| 67 | + TVM_FFI_EXTRA_CXX_API static TypeSchema FromJSON(const Any& obj); |
| 68 | + |
| 69 | + /*! |
| 70 | + * \brief Construct a TypeSchema from a JSON string. |
| 71 | + * \param json_str The JSON string to parse. |
| 72 | + * \return The constructed TypeSchema. |
| 73 | + */ |
| 74 | + TVM_FFI_EXTRA_CXX_API static TypeSchema FromJSONStr(const String& json_str); |
| 75 | + |
| 76 | + /*! |
| 77 | + * \brief Render a human-readable representation. |
| 78 | + * |
| 79 | + * Uses Python-style typing syntax: |
| 80 | + * - Unions: `"T1 | T2"` |
| 81 | + * - Optional: `"T | None"` |
| 82 | + * - Callables: `"Callable[[arg1, ...], ret]"` |
| 83 | + * - Containers: `"origin[arg1, ...]"` |
| 84 | + * |
| 85 | + * \return A human-readable string. |
| 86 | + */ |
| 87 | + TVM_FFI_EXTRA_CXX_API std::string Repr() const; |
| 88 | +}; |
| 89 | + |
| 90 | +/*! |
| 91 | + * \brief Pre-compiled type checker/converter built from a TypeSchema. |
| 92 | + * |
| 93 | + * On construction, the checker resolves the schema's origin string to an |
| 94 | + * efficient internal enum and pre-resolves any object type keys to runtime |
| 95 | + * type indices. This makes subsequent TryCast / CheckStrict calls fast |
| 96 | + * (no string comparisons on the hot path). |
| 97 | + * |
| 98 | + * The conversion semantics match the compile-time `TypeTraits<T>` system: |
| 99 | + * - `CheckStrict` ↔ `TypeTraits<T>::CheckAnyStrict` |
| 100 | + * - `TryCast` ↔ `TypeTraits<T>::TryCastFromAnyView` |
| 101 | + */ |
| 102 | +class TypeChecker { |
| 103 | + public: |
| 104 | + /*! |
| 105 | + * \brief Construct a TypeChecker from a TypeSchema. |
| 106 | + * \param schema The type schema to compile. |
| 107 | + */ |
| 108 | + TVM_FFI_EXTRA_CXX_API explicit TypeChecker(TypeSchema schema); |
| 109 | + |
| 110 | + /*! |
| 111 | + * \brief Try to convert src to the target type. |
| 112 | + * |
| 113 | + * May apply type coercion (e.g. int→float, List→Array with element |
| 114 | + * conversion). Returns std::nullopt if conversion is not possible. |
| 115 | + * |
| 116 | + * \param src The input value. |
| 117 | + * \return The converted value, or std::nullopt. |
| 118 | + */ |
| 119 | + TVM_FFI_EXTRA_CXX_API std::optional<Any> TryCast(AnyView src) const; |
| 120 | + |
| 121 | + /*! |
| 122 | + * \brief Convert src to the target type, or throw TypeError. |
| 123 | + * \param src The input value. |
| 124 | + * \return The converted value. |
| 125 | + */ |
| 126 | + TVM_FFI_EXTRA_CXX_API Any Cast(AnyView src) const; |
| 127 | + |
| 128 | + /*! |
| 129 | + * \brief Check if src strictly matches the target type (no conversion). |
| 130 | + * |
| 131 | + * For containers this recursively checks all elements. |
| 132 | + * |
| 133 | + * \param src The input value. |
| 134 | + * \return True if the value exactly matches the expected type. |
| 135 | + */ |
| 136 | + TVM_FFI_EXTRA_CXX_API bool CheckStrict(AnyView src) const; |
| 137 | + |
| 138 | + /*! \return The underlying TypeSchema. */ |
| 139 | + const TypeSchema& schema() const { return schema_; } |
| 140 | + |
| 141 | + private: |
| 142 | + /*! \brief Dispatch tag for efficient runtime type checking. */ |
| 143 | + enum class Kind : uint8_t { |
| 144 | + kAny, |
| 145 | + kNone, |
| 146 | + kInt, |
| 147 | + kBool, |
| 148 | + kFloat, |
| 149 | + kStr, |
| 150 | + kBytes, |
| 151 | + kDataType, |
| 152 | + kDevice, |
| 153 | + kOpaquePtr, |
| 154 | + kDLTensorPtr, |
| 155 | + kObject, |
| 156 | + kOptional, |
| 157 | + kArray, |
| 158 | + kList, |
| 159 | + kMap, |
| 160 | + kVariant, |
| 161 | + kTuple, |
| 162 | + kCallable, |
| 163 | + }; |
| 164 | + |
| 165 | + /*! \brief Resolve a TypeSchema into Kind + type index. */ |
| 166 | + static Kind ResolveKind(const std::string& origin, int32_t* resolved_type_index); |
| 167 | + |
| 168 | + Kind kind_; |
| 169 | + int32_t resolved_type_index_{-1}; |
| 170 | + TypeSchema schema_; |
| 171 | + std::vector<TypeChecker> args_; |
| 172 | +}; |
| 173 | + |
| 174 | +} // namespace ffi |
| 175 | +} // namespace tvm |
| 176 | + |
| 177 | +#endif // TVM_FFI_EXTRA_TYPE_CHECKER_H_ |
0 commit comments