return NULL;
 }
 
+/* Enable trace_probe - @flag must be TP_FLAG_TRACE or TP_FLAG_PROFILE */
+static int enable_trace_probe(struct trace_probe *tp, int flag)
+{
+       int ret = 0;
+
+       tp->flags |= flag;
+       if (tp->flags & (TP_FLAG_TRACE | TP_FLAG_PROFILE)) {
+               if (trace_probe_is_return(tp))
+                       ret = enable_kretprobe(&tp->rp);
+               else
+                       ret = enable_kprobe(&tp->rp.kp);
+       }
+
+       return ret;
+}
+
+/* Disable trace_probe - @flag must be TP_FLAG_TRACE or TP_FLAG_PROFILE */
+static void disable_trace_probe(struct trace_probe *tp, int flag)
+{
+       tp->flags &= ~flag;
+       if (!(tp->flags & (TP_FLAG_TRACE | TP_FLAG_PROFILE))) {
+               if (trace_probe_is_return(tp))
+                       disable_kretprobe(&tp->rp);
+               else
+                       disable_kprobe(&tp->rp.kp);
+       }
+}
+
 /* Unregister a trace_probe and probe_event: call with locking probe_lock */
 static void unregister_trace_probe(struct trace_probe *tp)
 {
        return TRACE_TYPE_PARTIAL_LINE;
 }
 
-static int probe_event_enable(struct ftrace_event_call *call)
-{
-       struct trace_probe *tp = (struct trace_probe *)call->data;
-
-       tp->flags |= TP_FLAG_TRACE;
-       if (trace_probe_is_return(tp))
-               return enable_kretprobe(&tp->rp);
-       else
-               return enable_kprobe(&tp->rp.kp);
-}
-
-static void probe_event_disable(struct ftrace_event_call *call)
-{
-       struct trace_probe *tp = (struct trace_probe *)call->data;
-
-       tp->flags &= ~TP_FLAG_TRACE;
-       if (!(tp->flags & (TP_FLAG_TRACE | TP_FLAG_PROFILE))) {
-               if (trace_probe_is_return(tp))
-                       disable_kretprobe(&tp->rp);
-               else
-                       disable_kprobe(&tp->rp.kp);
-       }
-}
-
 #undef DEFINE_FIELD
 #define DEFINE_FIELD(type, item, name, is_signed)                      \
        do {                                                            \
        head = this_cpu_ptr(call->perf_events);
        perf_trace_buf_submit(entry, size, rctx, entry->ret_ip, 1, regs, head);
 }
-
-static int probe_perf_enable(struct ftrace_event_call *call)
-{
-       struct trace_probe *tp = (struct trace_probe *)call->data;
-
-       tp->flags |= TP_FLAG_PROFILE;
-
-       if (trace_probe_is_return(tp))
-               return enable_kretprobe(&tp->rp);
-       else
-               return enable_kprobe(&tp->rp.kp);
-}
-
-static void probe_perf_disable(struct ftrace_event_call *call)
-{
-       struct trace_probe *tp = (struct trace_probe *)call->data;
-
-       tp->flags &= ~TP_FLAG_PROFILE;
-
-       if (!(tp->flags & TP_FLAG_TRACE)) {
-               if (trace_probe_is_return(tp))
-                       disable_kretprobe(&tp->rp);
-               else
-                       disable_kprobe(&tp->rp.kp);
-       }
-}
 #endif /* CONFIG_PERF_EVENTS */
 
 static __kprobes
 int kprobe_register(struct ftrace_event_call *event, enum trace_reg type)
 {
+       struct trace_probe *tp = (struct trace_probe *)event->data;
+
        switch (type) {
        case TRACE_REG_REGISTER:
-               return probe_event_enable(event);
+               return enable_trace_probe(tp, TP_FLAG_TRACE);
        case TRACE_REG_UNREGISTER:
-               probe_event_disable(event);
+               disable_trace_probe(tp, TP_FLAG_TRACE);
                return 0;
 
 #ifdef CONFIG_PERF_EVENTS
        case TRACE_REG_PERF_REGISTER:
-               return probe_perf_enable(event);
+               return enable_trace_probe(tp, TP_FLAG_PROFILE);
        case TRACE_REG_PERF_UNREGISTER:
-               probe_perf_disable(event);
+               disable_trace_probe(tp, TP_FLAG_PROFILE);
                return 0;
 #endif
        }
                        pr_warning("error on getting new probe.\n");
                        warn++;
                } else
-                       probe_event_enable(&tp->call);
+                       enable_trace_probe(tp, TP_FLAG_TRACE);
        }
 
        ret = command_trace_probe("r:testprobe2 kprobe_trace_selftest_target "
                        pr_warning("error on getting new probe.\n");
                        warn++;
                } else
-                       probe_event_enable(&tp->call);
+                       enable_trace_probe(tp, TP_FLAG_TRACE);
        }
 
        if (warn)