[k2] support covariance for future<T>

compiler/code-gen/files/type-tagger.cpp

@@ -4,6 +4,9 @@
 
 #include "compiler/code-gen/files/type-tagger.h"
 
+#include <set>
+#include <string>
+
 #include "common/algorithms/hashes.h"
 #include "compiler/code-gen/code-generator.h"
 #include "compiler/code-gen/common.h"
@@ -31,21 +34,42 @@ IncludesCollector TypeTagger::collect_includes() const noexcept {
 std::map<int, std::string> TypeTagger::collect_hash_of_types() const noexcept {
   // Be care, do not remove spaces from these types
   // TODO fix it?
-  std::set<std::string> sorted_types{
-    "bool",
-    "int64_t",
-    "Optional < int64_t >",
-    "void",
-    "thrown_exception",
-    "mixed",
-    "array< mixed >",
-    "Optional < string >",
-    "Optional < array< mixed > >",
-    "array< array< mixed > >",
-    "class_instance<C$KphpJobWorkerResponse>",
-    "class_instance<C$VK$TL$RpcResponse>",
-    "array< class_instance<C$VK$TL$RpcResponse> >",
-    "class_instance<C$PDOStatement>",
+
+  std::set<std::string> sorted_types{};
+  if (G->is_output_mode_k2()) {
+    sorted_types = std::set<std::string>{
+        "bool",
+        "int64_t",
+        "Optional < int64_t >",
+        "void",
+        "mixed",
+        "array< mixed >",
+        "Optional < string >",
+        "Optional < array< mixed > >",
+        "array< array< mixed > >",
+        "class_instance<C$KphpJobWorkerResponse>",
+        "class_instance<C$VK$TL$RpcResponse>",
+        "array< class_instance<C$VK$TL$RpcResponse> >",
+        // "class_instance<C$PDOStatement>",
+        // "thrown_exception",
+    };
+  } else {
+    sorted_types = std::set<std::string>{
+        "bool",
+        "int64_t",
+        "Optional < int64_t >",
+        "void",
+        "thrown_exception",
+        "mixed",
+        "array< mixed >",
+        "Optional < string >",
+        "Optional < array< mixed > >",
+        "array< array< mixed > >",
+        "class_instance<C$KphpJobWorkerResponse>",
+        "class_instance<C$VK$TL$RpcResponse>",
+        "array< class_instance<C$VK$TL$RpcResponse> >",
+        "class_instance<C$PDOStatement>",
+    };
   };
 
   for (const auto *type : forkable_types_) {
@@ -68,31 +92,55 @@ void TypeTagger::compile_tagger(CodeGenerator &W, const IncludesCollector &inclu
   W << ExternInclude{G->settings().runtime_headers.get()};
   W << includes << NL;
 
-  for (const auto &[hash, type] : hash_of_types) {
-    W << "template<>" << NL;
-    FunctionSignatureGenerator{W} << "int Storage::tagger<" << type << ">::get_tag() " << BEGIN;
-    W << "return " << hash << ";" << NL;
-    W << END << NL << NL;
+  for (const auto& [hash, type] : hash_of_types) {
+    if (G->is_output_mode_k2()) {
+      W << "template<>" << NL;
+      FunctionSignatureGenerator{W} << "int32_t kphp::forks::details::storage::tagger<" << type << ">::get_tag() " << BEGIN;
+      W << "return " << hash << SemicolonAndNL{};
+      W << END << NL << NL;
+    } else {
+      W << "template<>" << NL;
+      FunctionSignatureGenerator{W} << "int Storage::tagger<" << type << ">::get_tag() " << BEGIN;
+      W << "return " << hash << ";" << NL;
+      W << END << NL << NL;
+    }
   }
 
   W << CloseFile{};
 }
 
-void TypeTagger::compile_loader_header(CodeGenerator &W, const IncludesCollector &includes, const std::map<int, std::string> &hash_of_types) const noexcept {
+void TypeTagger::compile_loader_header(CodeGenerator& W, const IncludesCollector& includes, const std::map<int, std::string>& hash_of_types) const noexcept {
   W << OpenFile{loader_file_};
   W << ExternInclude{G->settings().runtime_headers.get()};
   W << includes << NL;
 
-  W << "template<typename T>" << NL;
-  FunctionSignatureGenerator{W} << "typename Storage::loader<T>::loader_fun Storage::loader<T>::get_function(int tag)" << BEGIN;
-  W << "switch(tag)" << BEGIN;
-  for (const auto &[hash, type] : hash_of_types) {
-    W << "case " << hash << ":"
-      << " return Storage::load_implementation_helper<" << type << ", T>::load;" << NL;
+  if (G->is_output_mode_k2()) {
+    W << "template<typename T>" << NL;
+    W << "requires(!std::same_as<T, int32_t>)" << NL;
+    FunctionSignatureGenerator{W}
+        << "typename kphp::forks::details::storage::loader<T>::loader_function_type kphp::forks::details::storage::loader<T>::get_loader(int32_t tag) "
+        << BEGIN;
+    W << "switch(tag)" << BEGIN;
+    for (const auto& [hash, type] : hash_of_types) {
+      W << "case " << hash << ":"
+        << " return kphp::forks::details::storage::loader_impl<" << type << ", T>::loader_function;" << NL;
+    }
+    W << END << NL;
+    W << "kphp::log::assertion(false);" << NL;
+    W << "return nullptr;" << NL;
+    W << END << NL;
+  } else {
+    W << "template<typename T>" << NL;
+    FunctionSignatureGenerator{W} << "typename Storage::loader<T>::loader_fun Storage::loader<T>::get_function(int tag)" << BEGIN;
+    W << "switch(tag)" << BEGIN;
+    for (const auto& [hash, type] : hash_of_types) {
+      W << "case " << hash << ":"
+        << " return Storage::load_implementation_helper<" << type << ", T>::load;" << NL;
+    }
+    W << END << NL;
+    W << "php_assert(0);" << NL;
+    W << END << NL;
   }
-  W << END << NL;
-  W << "php_assert(0);" << NL;
-  W << END << NL;
 
   W << CloseFile{};
 }
@@ -103,8 +151,15 @@ static void compile_loader_instantiations_batch(CodeGenerator &W, const Includes
   W << ExternInclude{G->settings().runtime_headers.get()};
   W << includes << NL;
 
-  for (auto it = begin; it != end; ++it) {
-    W << "template Storage::loader<" << *it << ">::loader_fun Storage::loader<" << *it << ">::get_function(int);" << NL;
+  if (G->is_output_mode_k2()) {
+    for (auto it = begin; it != end; ++it) {
+      W << "template kphp::forks::details::storage::loader<" << *it << ">::loader_function_type kphp::forks::details::storage::loader<" << *it
+        << ">::get_loader(int32_t);" << NL;
+    }
+  } else {
+    for (auto it = begin; it != end; ++it) {
+      W << "template Storage::loader<" << *it << ">::loader_fun Storage::loader<" << *it << ">::get_function(int);" << NL;
+    }
   }
 
   W << CloseFile{};

compiler/pipes/code-gen.cpp

@@ -117,10 +117,8 @@ void CodeGenF::on_finish(DataStream<std::unique_ptr<CodeGenRootCmd>> &os) {
 
   // TODO: should be done in lib mode also, but in some other way
   if (!G->is_output_mode_lib() && !G->is_output_mode_k2_lib()) {
-    if (!G->is_output_mode_k2()) {
-      code_gen_start_root_task(os, std::make_unique<TypeTagger>(vk::singleton<ForkableTypeStorage>::get().flush_forkable_types(),
-                                                                vk::singleton<ForkableTypeStorage>::get().flush_waitable_types()));
-    }
+    code_gen_start_root_task(os, std::make_unique<TypeTagger>(vk::singleton<ForkableTypeStorage>::get().flush_forkable_types(),
+                                                              vk::singleton<ForkableTypeStorage>::get().flush_waitable_types()));
     code_gen_start_root_task(os, std::make_unique<ShapeKeys>(TypeHintShape::get_all_registered_keys()));
     code_gen_start_root_task(os, std::make_unique<JsonEncoderTags>(std::move(all_json_encoders)));
     code_gen_start_root_task(os, std::make_unique<CmakeListsTxt>());

compiler/pipes/collect-forkable-types.cpp

@@ -9,14 +9,17 @@
 
 VertexPtr CollectForkableTypesPass::on_enter_vertex(VertexPtr root) {
   if (auto call = root.try_as<op_func_call>()) {
-    if (call->func_id->is_resumable) {
+    if (call->func_id->is_resumable || call->func_id->is_k2_fork || call->func_id->is_interruptible) {
       if (call->str_val == "wait") {
         waitable_types_.push_back(tinf::get_type(root));
       } else if (call->str_val == "wait_multi") {
         forkable_types_.push_back(tinf::get_type(root)->lookup_at_any_key());
         waitable_types_.push_back(tinf::get_type(root)->lookup_at_any_key());
       }
-      forkable_types_.push_back(tinf::get_type(root));
+
+      if (call->func_id->is_resumable || call->func_id->is_k2_fork) {
+        forkable_types_.push_back(tinf::get_type(root));
+      }
     } else if (call->str_val == "wait_synchronously") {
       waitable_types_.push_back(tinf::get_type(root));
     }

runtime-light/stdlib/fork/fork-functions.h

@@ -20,6 +20,7 @@
 #include "runtime-light/coroutine/type-traits.h"
 #include "runtime-light/stdlib/diagnostics/logs.h"
 #include "runtime-light/stdlib/fork/fork-state.h"
+#include "runtime-light/stdlib/fork/fork-storage.h"
 
 namespace kphp::forks {
 
@@ -82,7 +83,7 @@ auto wait(int64_t fork_id, std::chrono::nanoseconds timeout) noexcept -> kphp::c
   // Important: capture current fork's info pre-co_await.
   // Fork ID is not automatically preserved across suspension points
   auto current_fork_info{fork_instance_st.current_info()};
-  auto expected{co_await kphp::coro::io_scheduler::get().schedule(static_cast<kphp::coro::shared_task<return_type>>(std::move(fork_task)),
+  auto expected{co_await kphp::coro::io_scheduler::get().schedule(static_cast<kphp::coro::shared_task<kphp::forks::details::storage>>(std::move(fork_task)),
                                                                   detail::normalize_timeout(timeout))};
 
   // Execute essential housekeeping tasks to maintain proper state management.
@@ -95,7 +96,8 @@ auto wait(int64_t fork_id, std::chrono::nanoseconds timeout) noexcept -> kphp::c
   if (!expected) [[unlikely]] {
     co_return std::nullopt;
   }
-  co_return *std::move(expected);
+
+  co_return (*std::move(expected)).template load<return_type>();
 }
 
 } // namespace kphp::forks
@@ -104,7 +106,7 @@ template<typename T>
 requires(is_optional<T>::value || std::same_as<T, mixed> || is_class_instance<T>::value)
 kphp::coro::task<T> f$wait(int64_t fork_id, double timeout = -1.0) noexcept {
   auto opt_result{co_await kphp::forks::id_managed(
-      kphp::forks::wait<internal_optional_type_t<T>>(fork_id, std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::duration<double>{timeout})))};
+      kphp::forks::wait<T>(fork_id, std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::duration<double>{timeout})))};
   co_return opt_result ? T{*std::move(opt_result)} : T{};
 }
 

runtime-light/stdlib/fork/fork-state.h

@@ -16,6 +16,7 @@
 #include "runtime-light/coroutine/task.h"
 #include "runtime-light/stdlib/diagnostics/exception-types.h"
 #include "runtime-light/stdlib/diagnostics/logs.h"
+#include "runtime-light/stdlib/fork/fork-storage.h"
 
 namespace kphp::forks {
 
@@ -30,7 +31,7 @@ struct ForkInstanceState final : private vk::not_copyable {
   struct fork_info final {
     bool awaited{};
     Throwable thrown_exception;
-    std::optional<kphp::coro::shared_task<>> opt_handle;
+    std::optional<kphp::coro::shared_task<kphp::forks::details::storage>> opt_handle;
   };
 
 private:
@@ -48,15 +49,26 @@ struct ForkInstanceState final : private vk::not_copyable {
   static ForkInstanceState& get() noexcept;
 
   template<typename return_type>
-  std::pair<int64_t, kphp::coro::shared_task<return_type>> create_fork(kphp::coro::task<return_type> task) noexcept {
-    static constexpr auto fork_coroutine{[](kphp::coro::task<return_type> task, int64_t fork_id) noexcept -> kphp::coro::shared_task<return_type> {
-      ForkInstanceState::get().current_id = fork_id;
-      co_return co_await std::move(task);
-    }};
+  std::pair<int64_t, kphp::coro::shared_task<kphp::forks::details::storage>> create_fork(kphp::coro::task<return_type> task) noexcept {
+    static constexpr auto fork_coroutine{
+        [](kphp::coro::task<return_type> task, int64_t fork_id) noexcept -> kphp::coro::shared_task<kphp::forks::details::storage> {
+          ForkInstanceState::get().current_id = fork_id;
+
+          kphp::forks::details::storage s{};
+          if constexpr (std::same_as<return_type, void>) {
+            co_await std::move(task);
+            s.store();
+          } else {
+            s.store<return_type>(co_await std::move(task));
+          }
+          co_return s;
+        }};
 
     const int64_t fork_id{next_fork_id++};
     auto fork_task{std::invoke(fork_coroutine, std::move(task), fork_id)};
-    forks.emplace(fork_id, fork_info{.awaited = {}, .thrown_exception = {}, .opt_handle = static_cast<kphp::coro::shared_task<>>(fork_task)});
+    forks.emplace(
+        fork_id,
+        fork_info{.awaited = {}, .thrown_exception = {}, .opt_handle = static_cast<kphp::coro::shared_task<kphp::forks::details::storage>>(fork_task)});
     return std::make_pair(fork_id, std::move(fork_task));
   }
 

runtime-light/stdlib/fork/fork-storage.cpp

@@ -0,0 +1,18 @@
+// Compiler for PHP (aka KPHP)
+// Copyright (c) 2025 LLC «V Kontakte»
+// Distributed under the GPL v3 License, see LICENSE.notice.txt
+
+#include "runtime-light/stdlib/fork/fork-storage.h"
+
+#include <optional>
+
+#include "runtime-light/stdlib/diagnostics/logs.h"
+
+namespace kphp::forks::details {
+
+auto storage::store() noexcept -> void {
+  kphp::log::assertion(!m_opt_tag.has_value());
+  m_opt_tag.emplace(tagger<void>::get_tag());
+}
+
+} // namespace kphp::forks::details

runtime-light/stdlib/fork/fork-storage.h

@@ -0,0 +1,101 @@
+// Compiler for PHP (aka KPHP)
+// Copyright (c) 2025 LLC «V Kontakte»
+// Distributed under the GPL v3 License, see LICENSE.notice.txt
+
+#pragma once
+
+#include <concepts>
+#include <cstdint>
+#include <optional>
+#include <type_traits>
+#include <utility>
+
+#include "runtime-common/core/runtime-core.h"
+#include "runtime-common/core/utils/small-object-storage.h"
+#include "runtime-light/stdlib/diagnostics/logs.h"
+
+namespace kphp::forks::details {
+
+class storage {
+  using data_type = small_object_storage<sizeof(mixed)>;
+
+  template<typename From, typename To, typename Convertible = std::is_convertible<From, To>::type>
+  struct loader_impl;
+
+  template<typename From, typename To>
+  requires(!std::same_as<From, int32_t> && !std::same_as<To, int32_t>)
+  struct loader_impl<From, To, std::true_type> {
+    static To loader_function(data_type& storage) noexcept {
+      From* data{storage.get<From>()};
+      To result(std::move(*data));
+      storage.destroy<From>();
+      return result;
+    }
+  };
+
+  template<typename From, typename To>
+  requires(!std::same_as<From, int32_t> && !std::same_as<To, int32_t>)
+  struct loader_impl<From, To, std::false_type> {
+    static To loader_function([[maybe_unused]] data_type& storage) noexcept {
+      kphp::log::assertion(false);
+      return To();
+    }
+  };
+
+  template<>
+  struct loader_impl<void, void, std::true_type> {
+    static void loader_function([[maybe_unused]] data_type& storage) noexcept {}
+  };
+
+  template<typename From>
+  requires(!std::same_as<From, int32_t>)
+  struct loader_impl<From, void, std::false_type> {
+    static void loader_function(data_type& storage) noexcept {
+      loader_impl<From, From>::loader_function(storage);
+    }
+  };
+
+  std::optional<int32_t> m_opt_tag;
+  data_type m_data;
+
+public:
+  template<typename T>
+  requires(!std::same_as<T, int32_t>)
+  struct tagger {
+    static auto get_tag() noexcept -> int32_t;
+  };
+
+  template<typename T>
+  requires(!std::same_as<T, int32_t>)
+  struct loader {
+    using loader_function_type = T (*)(data_type&);
+
+    static auto get_loader(int32_t tag) noexcept -> loader_function_type;
+  };
+
+  // `std::type_identity_t` is used to disable type deduction for `store` function.
+  // It should be called with exactly the same type parameter as `load` function.
+  // So, to prevent a bug we decided to forbid type deduction here.
+  template<typename T>
+  requires(!std::same_as<T, int32_t> && !std::same_as<T, void>)
+  auto store(std::type_identity_t<T> val) noexcept -> void {
+    kphp::log::assertion(!m_opt_tag.has_value());
+    m_opt_tag.emplace(tagger<T>::get_tag());
+    m_data.emplace<T>(std::move(val));
+  }
+
+  auto store() noexcept -> void;
+
+  template<typename T>
+  requires(!std::same_as<T, int32_t>)
+  auto load() noexcept -> T {
+    kphp::log::assertion(m_opt_tag.has_value());
+    const auto tag{*m_opt_tag};
+    m_opt_tag = std::nullopt;
+    if constexpr (!std::same_as<T, void>) {
+      return loader<T>::get_loader(tag)(m_data);
+    }
+  }
+};
+
+} // namespace kphp::forks::details