From 09ed245281300cfb84969614ab6e0a0170fd8f84 Mon Sep 17 00:00:00 2001
From: Jamieson Pryor <jwhpryor@google.com>
Date: Tue, 9 Jul 2024 15:14:41 -0700
Subject: [PATCH] Fixes field `Static`s that were broken when using `Self`.

https://github.com/google/jni-bind/issues/311

PiperOrigin-RevId: 650772628
---
 implementation/field_ref.h        |  8 ++++++-
 implementation/static_ref_test.cc | 37 +++++++++++++++++++++++++++----
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/implementation/field_ref.h b/implementation/field_ref.h
index 24c9b9ff..05e1ba51 100644
--- a/implementation/field_ref.h
+++ b/implementation/field_ref.h
@@ -20,6 +20,7 @@
 #include <mutex>
 #include <string_view>
 #include <tuple>
+#include <type_traits>
 #include <utility>
 
 #include "implementation/class_ref.h"
@@ -57,6 +58,8 @@ class FieldRef {
   using IdT = Id<JniT, field_type, I>;
   using FieldSelectionT = FieldSelection<JniT, I>;
 
+  using SelfIdT = typename IdT::template ChangeIdType<IdType::CLASS>;
+
   explicit FieldRef(jclass class_ref, jobject object_ref)
       : class_ref_(class_ref), object_ref_(object_ref) {}
 
@@ -83,7 +86,10 @@ class FieldRef {
     });
   }
 
-  using ReturnProxied = Return_t<typename IdT::MaterializeCDeclT, IdT>;
+  using ReturnProxied =
+      std::conditional_t<IdT::kIsSelf,
+                         Return_t<typename SelfIdT::MaterializeCDeclT, SelfIdT>,
+                         Return_t<typename IdT::MaterializeCDeclT, IdT> >;
 
   const auto& SelfVal() {
     if constexpr (IdT::kIsStatic) {
diff --git a/implementation/static_ref_test.cc b/implementation/static_ref_test.cc
index 0819cda2..b592ff66 100644
--- a/implementation/static_ref_test.cc
+++ b/implementation/static_ref_test.cc
@@ -31,6 +31,7 @@ using ::jni::LocalString;
 using ::jni::Method;
 using ::jni::Params;
 using ::jni::Rank;
+using ::jni::Self;
 using ::jni::Static;
 using ::jni::StaticRef;
 using ::jni::test::JniTest;
@@ -57,7 +58,8 @@ static constexpr Class kClass{
         Field{"floatField", jfloat{}},
         Field{"doubleField", jdouble{}},
         Field{"stringField", jstring{}},
-        Field{"classField", Class{"kClass2"}}
+        Field{"classField", Class{"kClass2"}},
+        Field{"selfField", Self{}}
       },
 };
 // clang-format on
@@ -190,6 +192,25 @@ TEST_F(JniTest, StaticField_ObjectSet) {
       LocalObject<kClass2>{AdoptLocal{}, Fake<jobject>()});
 }
 
+TEST_F(JniTest, StaticField_SelfGet) {
+  EXPECT_CALL(*env_, GetStaticFieldID(_, StrEq("selfField"), StrEq("LkClass;")))
+      .WillOnce(Return(Fake<jfieldID>()));
+  EXPECT_CALL(*env_, GetStaticObjectField(_, Fake<jfieldID>()))
+      .WillOnce(Return(Fake<jobject>()));
+
+  jni::LocalObject<kClass> obj = StaticRef<kClass>{}["selfField"].Get();
+}
+
+TEST_F(JniTest, StaticField_SelfSet) {
+  EXPECT_CALL(*env_, GetStaticFieldID(_, StrEq("selfField"), StrEq("LkClass;")))
+      .WillOnce(Return(Fake<jfieldID>()));
+  EXPECT_CALL(*env_,
+              SetStaticObjectField(_, Fake<jfieldID>(), Fake<jobject>()));
+
+  StaticRef<kClass>{}["selfField"].Set(
+      LocalObject<kClass>{AdoptLocal{}, Fake<jobject>()});
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // Static Methods.
 ////////////////////////////////////////////////////////////////////////////////
@@ -210,6 +231,7 @@ static constexpr Class kMethodClass{
         Method{"objectMethod", ::jni::Return{Class{"kClass2"}}},
         Method{"rank1ArrayMethod", ::jni::Return{Array{Class{"kClass2"}}}},
         Method{"rank2ArrayMethod", ::jni::Return{Array{Class{"kClass2"}, Rank<2>{}}}},
+        Method{"selfMethod", ::jni::Return{Self{}}},
 
         Method{"simpleFunc", ::jni::Return{int{}}, Params{jfloat{}}},
         Method{"complexFunc", ::jni::Return{float{}},
@@ -248,8 +270,10 @@ TEST_F(JniTest, StaticExerciseAllReturns) {
                                        StrEq("()[LkClass2;")));
   EXPECT_CALL(*env_, GetStaticMethodID(_, StrEq("rank2ArrayMethod"),
                                        StrEq("()[[LkClass2;")));
+  EXPECT_CALL(*env_, GetStaticMethodID(_, StrEq("selfMethod"),
+                                       StrEq("()LkMethodClass;")));
 
-  EXPECT_CALL(*env_, CallStaticObjectMethodV).Times(4);
+  EXPECT_CALL(*env_, CallStaticObjectMethodV).Times(5);
 
   StaticRef<kMethodClass>{}("booleanMethod");
   StaticRef<kMethodClass>{}("byteMethod");
@@ -261,10 +285,11 @@ TEST_F(JniTest, StaticExerciseAllReturns) {
   StaticRef<kMethodClass>{}("doubleMethod");
   StaticRef<kMethodClass>{}("stringMethod");
 
-  // It would be more complete to exercise all types here.
   StaticRef<kMethodClass>{}("objectMethod");
   StaticRef<kMethodClass>{}("rank1ArrayMethod");
   StaticRef<kMethodClass>{}("rank2ArrayMethod");
+
+  LocalObject<kClass> self_ret = StaticRef<kMethodClass>{}("selfMethod");
 }
 
 // clang-format off
@@ -283,6 +308,7 @@ static constexpr Class kMethodClassSingleParam{
         Method{"objectMethod", ::jni::Return<void>{}, Params{Class{"kClass2"}}},
         Method{"rank1ArrayMethod", ::jni::Return<void>{}, Params{Array{Class{"kClass2"}}}},
         Method{"rank2ArrayMethod", ::jni::Return<void>{}, Params{Array{Class{"kClass2"}, Rank<2>{}}}},
+        Method{"selfMethod", ::jni::Return<void>{}, Params{Self{}}}
       },
 };
 // clang-format on
@@ -307,8 +333,10 @@ TEST_F(JniTest, StaticExerciseAllTypesThroughSingleParam) {
                                        StrEq("([LkClass2;)V")));
   EXPECT_CALL(*env_, GetStaticMethodID(_, StrEq("rank2ArrayMethod"),
                                        StrEq("([[LkClass2;)V")));
+  EXPECT_CALL(*env_, GetStaticMethodID(_, StrEq("selfMethod"),
+                                       StrEq("(LkMethodClassSingleParam;)V")));
 
-  EXPECT_CALL(*env_, CallStaticVoidMethodV).Times(12);
+  EXPECT_CALL(*env_, CallStaticVoidMethodV).Times(13);
 
   StaticRef<kMethodClassSingleParam>{}("booleanMethod", jboolean{true});
   StaticRef<kMethodClassSingleParam>{}("byteMethod", jbyte{1});
@@ -326,6 +354,7 @@ TEST_F(JniTest, StaticExerciseAllTypesThroughSingleParam) {
                                        Fake<jobjectArray>());
   StaticRef<kMethodClassSingleParam>{}("rank2ArrayMethod",
                                        Fake<jobjectArray>());
+  StaticRef<kMethodClassSingleParam>{}("selfMethod", Fake<jobject>());
 }
 
 TEST_F(JniTest, StaticExerciseComplexSetOfParams) {