#include "trace.h"
 #include "trace_output.h"
 
+struct fgraph_data {
+       pid_t           last_pid;
+       int             depth;
+};
+
 #define TRACE_GRAPH_INDENT     2
 
 /* Flag options */
 
 /* If the pid changed since the last trace, output this event */
 static enum print_line_t
-verif_pid(struct trace_seq *s, pid_t pid, int cpu, pid_t *last_pids_cpu)
+verif_pid(struct trace_seq *s, pid_t pid, int cpu, struct fgraph_data *data)
 {
        pid_t prev_pid;
        pid_t *last_pid;
        int ret;
 
-       if (!last_pids_cpu)
+       if (!data)
                return TRACE_TYPE_HANDLED;
 
-       last_pid = per_cpu_ptr(last_pids_cpu, cpu);
+       last_pid = &(per_cpu_ptr(data, cpu)->last_pid);
 
        if (*last_pid == pid)
                return TRACE_TYPE_HANDLED;
                struct ftrace_graph_ent_entry *entry,
                struct ftrace_graph_ret_entry *ret_entry, struct trace_seq *s)
 {
+       struct fgraph_data *data = iter->private;
        struct ftrace_graph_ret *graph_ret;
        struct ftrace_graph_ent *call;
        unsigned long long duration;
        call = &entry->graph_ent;
        duration = graph_ret->rettime - graph_ret->calltime;
 
+       if (data) {
+               int cpu = iter->cpu;
+               int *depth = &(per_cpu_ptr(data, cpu)->depth);
+
+               /*
+                * Comments display at + 1 to depth. Since
+                * this is a leaf function, keep the comments
+                * equal to this depth.
+                */
+               *depth = call->depth - 1;
+       }
+
        /* Overhead */
        ret = print_graph_overhead(duration, s);
        if (!ret)
 }
 
 static enum print_line_t
-print_graph_entry_nested(struct ftrace_graph_ent_entry *entry,
-                       struct trace_seq *s, pid_t pid, int cpu)
+print_graph_entry_nested(struct trace_iterator *iter,
+                        struct ftrace_graph_ent_entry *entry,
+                        struct trace_seq *s, int cpu)
 {
-       int i;
-       int ret;
        struct ftrace_graph_ent *call = &entry->graph_ent;
+       struct fgraph_data *data = iter->private;
+       int ret;
+       int i;
+
+       if (data) {
+               int cpu = iter->cpu;
+               int *depth = &(per_cpu_ptr(data, cpu)->depth);
+
+               *depth = call->depth;
+       }
 
        /* No overhead */
        ret = print_graph_overhead(-1, s);
 print_graph_prologue(struct trace_iterator *iter, struct trace_seq *s,
                     int type, unsigned long addr)
 {
+       struct fgraph_data *data = iter->private;
        struct trace_entry *ent = iter->ent;
-       pid_t *last_pid = iter->private;
        int cpu = iter->cpu;
        int ret;
 
        /* Pid */
-       if (verif_pid(s, ent->pid, cpu, last_pid) == TRACE_TYPE_PARTIAL_LINE)
+       if (verif_pid(s, ent->pid, cpu, data) == TRACE_TYPE_PARTIAL_LINE)
                return TRACE_TYPE_PARTIAL_LINE;
 
        if (type) {
        if (leaf_ret)
                return print_graph_entry_leaf(iter, field, leaf_ret, s);
        else
-               return print_graph_entry_nested(field, s, iter->ent->pid, cpu);
+               return print_graph_entry_nested(iter, field, s, cpu);
 
 }
 
 print_graph_return(struct ftrace_graph_ret *trace, struct trace_seq *s,
                   struct trace_entry *ent, struct trace_iterator *iter)
 {
-       int i;
-       int ret;
-       int cpu = iter->cpu;
-       pid_t pid = ent->pid;
        unsigned long long duration = trace->rettime - trace->calltime;
+       struct fgraph_data *data = iter->private;
+       pid_t pid = ent->pid;
+       int cpu = iter->cpu;
+       int ret;
+       int i;
+
+       if (data) {
+               int cpu = iter->cpu;
+               int *depth = &(per_cpu_ptr(data, cpu)->depth);
+
+               /*
+                * Comments display at + 1 to depth. This is the
+                * return from a function, we now want the comments
+                * to display at the same level of the bracket.
+                */
+               *depth = trace->depth - 1;
+       }
 
        if (print_graph_prologue(iter, s, 0, 0))
                return TRACE_TYPE_PARTIAL_LINE;
 print_graph_comment(struct bprint_entry *trace, struct trace_seq *s,
                   struct trace_entry *ent, struct trace_iterator *iter)
 {
-       int i;
+       struct fgraph_data *data = iter->private;
+       int depth = 0;
        int ret;
+       int i;
+
+       if (data)
+               depth = per_cpu_ptr(data, iter->cpu)->depth;
 
        if (print_graph_prologue(iter, s, 0, 0))
                return TRACE_TYPE_PARTIAL_LINE;
        }
 
        /* Indentation */
-       if (trace->depth > 0)
-               for (i = 0; i < (trace->depth + 1) * TRACE_GRAPH_INDENT; i++) {
+       if (depth > 0)
+               for (i = 0; i < (depth + 1) * TRACE_GRAPH_INDENT; i++) {
                        ret = trace_seq_printf(s, " ");
                        if (!ret)
                                return TRACE_TYPE_PARTIAL_LINE;
 
 static void graph_trace_open(struct trace_iterator *iter)
 {
-       /* pid on the last trace processed */
-       pid_t *last_pid = alloc_percpu(pid_t);
+       /* pid and depth on the last trace processed */
+       struct fgraph_data *data = alloc_percpu(struct fgraph_data);
        int cpu;
 
-       if (!last_pid)
+       if (!data)
                pr_warning("function graph tracer: not enough memory\n");
        else
                for_each_possible_cpu(cpu) {
-                       pid_t *pid = per_cpu_ptr(last_pid, cpu);
+                       pid_t *pid = &(per_cpu_ptr(data, cpu)->last_pid);
+                       int *depth = &(per_cpu_ptr(data, cpu)->depth);
                        *pid = -1;
+                       *depth = 0;
                }
 
-       iter->private = last_pid;
+       iter->private = data;
 }
 
 static void graph_trace_close(struct trace_iterator *iter)