From b950d6f8a9e4c33f098cb62fce77b083aa4377ee Mon Sep 17 00:00:00 2001
From: Sasha Goldshtein <goldshtn@gmail.com>
Date: Mon, 21 Mar 2016 04:06:15 -0700
Subject: [PATCH] Moved auto-includes helper to __init__.py

---
 src/python/bcc/__init__.py | 24 ++++++++++++++++++++++++
 tools/argdist.py           | 20 +-------------------
 tools/trace.py             | 21 +--------------------
 3 files changed, 26 insertions(+), 39 deletions(-)

diff --git a/src/python/bcc/__init__.py b/src/python/bcc/__init__.py
index 452997d..574db37 100644
--- a/src/python/bcc/__init__.py
+++ b/src/python/bcc/__init__.py
@@ -73,6 +73,30 @@ class BPF(object):
     _lib_load_address_cache = {}
     _lib_symbol_cache = {}
 
+    _auto_includes = {
+        "linux/time.h"      : ["time"],
+        "linux/fs.h"        : ["fs", "file"],
+        "linux/blkdev.h"    : ["bio", "request"],
+        "linux/slab.h"      : ["alloc"],
+        "linux/netdevice.h" : ["sk_buff", "net_device"]
+    }
+
+    @classmethod
+    def generate_auto_includes(cls, program_words):
+        """
+        Generates #include statements automatically based on a set of
+        recognized types such as sk_buff and bio. The input is all the words
+        that appear in the BPF program, and the output is a (possibly empty)
+        string of #include statements, such as "#include <linux/fs.h>".
+        """
+        headers = ""
+        for header, keywords in cls._auto_includes.items():
+            for keyword in keywords:
+                for word in program_words:
+                    if keyword in word and header not in headers:
+                        headers += "#include <%s>\n" % header
+        return headers
+
     # defined for compatibility reasons, to be removed
     Table = Table
 
diff --git a/tools/argdist.py b/tools/argdist.py
index 738a717..8f8327d 100755
--- a/tools/argdist.py
+++ b/tools/argdist.py
@@ -36,24 +36,6 @@ int PROBENAME(struct pt_regs *ctx SIGNATURE)
 """
         next_probe_index = 0
         aliases = { "$PID": "bpf_get_current_pid_tgid()" }
-        auto_includes = {
-                "linux/time.h"      : ["time"],
-                "linux/fs.h"        : ["fs", "file"],
-                "linux/blkdev.h"    : ["bio", "request"],
-                "linux/slab.h"      : ["alloc"],
-                "linux/netdevice.h" : ["sk_buff", "net_device"]
-        }
-
-        @staticmethod
-        def generate_auto_includes(specifiers):
-                headers = ""
-                for header, keywords in Specifier.auto_includes.items():
-                        for keyword in keywords:
-                                for specifier in specifiers:
-                                        if keyword in specifier:
-                                                headers += "#include <%s>\n" \
-                                                           % header
-                return headers
 
         def _substitute_aliases(self, expr):
                 if expr is None:
@@ -590,7 +572,7 @@ struct __string_t { char s[%d]; };
                 """ % self.args.string_size
                 for include in (self.args.include or []):
                         bpf_source += "#include <%s>\n" % include
-                bpf_source += Specifier.generate_auto_includes(
+                bpf_source += BPF.generate_auto_includes(
                                 map(lambda s: s.raw_spec, self.specifiers))
                 bpf_source += Tracepoint.generate_decl()
                 bpf_source += Tracepoint.generate_entry_probe()
diff --git a/tools/trace.py b/tools/trace.py
index b159510..33d0afa 100755
--- a/tools/trace.py
+++ b/tools/trace.py
@@ -328,25 +328,6 @@ int %s(struct pt_regs *ctx)
         def _time_off_str(cls, timestamp_ns):
                 return "%.6f" % (1e-9 * (timestamp_ns - cls.first_ts))
 
-        auto_includes = {
-                "linux/time.h"      : ["time"],
-                "linux/fs.h"        : ["fs", "file"],
-                "linux/blkdev.h"    : ["bio", "request"],
-                "linux/slab.h"      : ["alloc"],
-                "linux/netdevice.h" : ["sk_buff"]
-        }
-
-        @classmethod
-        def generate_auto_includes(cls, probes):
-                headers = ""
-                for header, keywords in cls.auto_includes.items():
-                        for keyword in keywords:
-                                for probe in probes:
-                                        if keyword in probe:
-                                                headers += "#include <%s>\n" \
-                                                           % header
-                return headers
-
         def _display_function(self):
                 if self.probe_type != 't':
                         return self.function
@@ -468,7 +449,7 @@ trace 't:block:block_rq_complete "sectors=%d", tp.nr_sector'
 #include <linux/sched.h>        /* For TASK_COMM_LEN */
 
 """
-                self.program += Probe.generate_auto_includes(
+                self.program += BPF.generate_auto_includes(
                         map(lambda p: p.raw_probe, self.probes))
                 self.program += Tracepoint.generate_decl()
                 self.program += Tracepoint.generate_entry_probe()
-- 
2.7.4