ext/cs__common/cs__common.c in contrast-agent-6.0.0 vs ext/cs__common/cs__common.c in contrast-agent-6.1.0

@@ -57,40 +57,69 @@
 
 VALUE contrast_patcher() {
     return patcher;
 }
 
+/* register instance alias patch */
 VALUE contrast_register_patch(const char *module_name, const char *method_name,
                               VALUE(c_fn)(const int, VALUE *, const VALUE)) {
     return _contrast_register_patch(module_name, method_name, c_fn,
                                     IMPL_ALIAS_INSTANCE);
 }
 
+/* register singleton alias patch */
 VALUE contrast_register_singleton_patch(const char *module_name,
                                         const char *method_name,
                                         VALUE(c_fn)(const int, VALUE *,
                                                     const VALUE)) {
     return _contrast_register_patch(module_name, method_name, c_fn,
                                     IMPL_ALIAS_SINGLETON);
 }
 
+/* register instance prepend patch */
 VALUE contrast_register_prepend_patch(const char *module_name,
                                       const char *method_name,
                                       VALUE(c_fn)(const int, VALUE *,
                                                   const VALUE)) {
     return _contrast_register_patch(module_name, method_name, c_fn,
                                     IMPL_PREPEND_INSTANCE);
 }
 
+/* register singleton prepend patch */
 VALUE contrast_register_singleton_prepend_patch(const char *module_name,
                                                 const char *method_name,
                                                 VALUE(c_fn)(const int, VALUE *,
                                                             const VALUE)) {
     return _contrast_register_patch(module_name, method_name, c_fn,
                                     IMPL_PREPEND_SINGLETON);
 }
 
+/* check if method is prepended and register instance alias or prepend patch */
+/* module name c_char "Module"; */
+/* method name c_char  "method"; */
+/* c_func => pointer */
+VALUE contrast_check_and_register_instance_patch(
+    const char *module_name, const char *method_name,
+    VALUE(c_fn)(const int, VALUE *, const VALUE)) {
+
+    VALUE object, method, is_prepended, patch_type;
+    /* check if method is prepended */
+    object = rb_const_get(rb_cObject, rb_intern(module_name));
+    method = ID2SYM(rb_intern(method_name));
+    is_prepended = contrast_check_prepended(object, method, Qtrue);
+
+    if (is_prepended == Qtrue) {
+        /* prepend patch */
+        return _contrast_register_patch(module_name, method_name, c_fn,
+                                        IMPL_PREPEND_INSTANCE);
+    } else {
+        /* alias patch */
+        return _contrast_register_patch(module_name, method_name, c_fn,
+                                        IMPL_ALIAS_INSTANCE);
+    }
+}
+
 static VALUE
 _contrast_register_patch(const char *module_name, const char *method_name,
                          VALUE(c_fn)(const int, VALUE *, const VALUE),
                          patch_impl patch) {
     VALUE contrast_bind_module = rb_funcall(rb_cModule, rb_intern("new"), 0);
@@ -131,10 +160,11 @@
     case IMPL_ALIAS_SINGLETON:
         impl = ID2SYM(rb_sym_alias_singleton);
         break;
     case IMPL_PREPEND_INSTANCE:
         impl = ID2SYM(rb_sym_prepend_instance);
+        break;
     case IMPL_PREPEND_SINGLETON:
         impl = ID2SYM(rb_sym_prepend_singleton);
         break;
     }
 
@@ -149,10 +179,75 @@
         FIX2INT(rb_funcall(rb_const_get(rb_cObject, rb_intern("RUBY_VERSION")),
                            rb_intern("to_i"), 0));
     return ruby_version < 3;
 }
 
+/* used for direct check on object: String.cs__prepended? *args  */
+extern VALUE contrast_check_prepended(VALUE self, VALUE method_name,
+                                      VALUE is_instance) {
+    return _contrast_check_prepended(self, method_name, is_instance);
+}
+
+/* used for passing object to look if not called on itself.
+Contrast::Agent::Assess.cs__object_method_prepended? object, :method_name,
+true/false */
+extern VALUE contrast_lookout_prepended(VALUE self, VALUE object_name,
+                                        VALUE method_name, VALUE is_instance) {
+    /* object_name must be the object, the self value is needed to prevent
+     lookout for self, since is always passed first we skip it */
+    VALUE result =
+        _contrast_check_prepended(object_name, method_name, is_instance);
+    return result;
+}
+
+static VALUE _contrast_check_prepended(VALUE object, VALUE method_name,
+                                       VALUE is_instance) {
+    VALUE entry, ancestors, object_idx, entry_methods;
+    VALUE result = Qfalse;
+    int i;
+    int y;
+
+    /* get self ancestors */
+    ancestors = rb_mod_ancestors(object);
+    /* get the size of the array */
+    int length = RARRAY_LEN(ancestors);
+    /* Locate self in ancestors: */
+    for (i = 0; i < length; ++i) {
+        entry = rb_ary_entry(ancestors, i);
+        if (entry == object) {
+            object_idx = i;
+            break;
+        }
+    }
+
+    /* find all the prepended modules */
+    /* we have the object place in ancestors: */
+    /* [suspect, suspect, object, ...] */
+    for (i = 0; i < object_idx; ++i) {
+        entry = rb_ary_entry(ancestors, i);
+        if (is_instance == Qtrue) {
+            entry_methods = rb_class_instance_methods(1, entry, entry);
+        } else {
+            entry_methods = rb_obj_singleton_methods(1, entry, entry);
+        }
+
+        /* Loop through the instance/singleton methods of the prepended modules
+         */
+        int entry_methods_length = RARRAY_LEN(entry_methods);
+        for (y = 0; y <= entry_methods_length; ++y) {
+            if (rb_ary_entry(entry_methods, y) == method_name) {
+                result = Qtrue;
+                break;
+            }
+        }
+        if (result == Qtrue) {
+            break;
+        }
+    }
+    return result;
+}
+
 void Init_cs__common(void) {
     cs__send_method = rb_intern("send");
     cs__alias_method_sym = ID2SYM(rb_intern("alias_method"));
 
     /* Define symbols */
@@ -189,6 +284,12 @@
     assess_policy = rb_define_module_under(assess, "Policy");
     assess_propagator = rb_define_module_under(assess_policy, "Propagator");
 
     core_extensions = rb_define_module_under(contrast, "Extension");
     core_assess = rb_define_module_under(core_extensions, "Assess");
+    /* defined for direct object check */
+    rb_define_singleton_method(rb_cObject, "cs__prepended?",
+                               contrast_check_prepended, 2);
+    /* defined for object lookout */
+    rb_define_singleton_method(assess, "cs__object_method_prepended?",
+                               contrast_lookout_prepended, 4);
 }