Merge pull request #188 from iovisor/probe_short_name

probe type short names 

src/ast/semantic_analyser.cpp

@@ -683,6 +683,8 @@ void SemanticAnalyser::visit(Predicate &pred)
 
 void SemanticAnalyser::visit(AttachPoint &ap)
 {
+  ap.provider = probetypeName(ap.provider);
+
   if (ap.provider == "kprobe" || ap.provider == "kretprobe") {
     if (ap.target != "")
       err_ << "kprobes should not have a target" << std::endl;

src/types.cpp

@@ -55,33 +55,24 @@ std::string typestr(Type t)
 
 ProbeType probetype(const std::string &type)
 {
-  if (type == "kprobe")
-    return ProbeType::kprobe;
-  else if (type == "kretprobe")
-    return ProbeType::kretprobe;
-  else if (type == "uprobe")
-    return ProbeType::uprobe;
-  else if (type == "uretprobe")
-    return ProbeType::uretprobe;
-  else if (type == "usdt")
-    return ProbeType::usdt;
-  else if (type == "BEGIN")
-    return ProbeType::uprobe;
-  else if (type == "END")
-    return ProbeType::uprobe;
-  else if (type == "tracepoint")
-    return ProbeType::tracepoint;
-  else if (type == "profile")
-    return ProbeType::profile;
-  else if (type == "interval")
-    return ProbeType::interval;
-  else if (type == "software")
-    return ProbeType::software;
-  else if (type == "hardware")
-    return ProbeType::hardware;
+  for (int i = 0; i < sizeof(PROBE_LIST); i++)
+  {
+    if (type == PROBE_LIST[i].name || type == PROBE_LIST[i].abbr)
+      return PROBE_LIST[i].type;
+  }
   abort();
 }
 
+std::string probetypeName(const std::string &type)
+{
+  for (int i = 0; i < sizeof(PROBE_LIST); i++)
+  {
+    if (type == PROBE_LIST[i].name || type == PROBE_LIST[i].abbr)
+      return PROBE_LIST[i].name;
+  }
+  return type;
+}
+
 uint64_t asyncactionint(AsyncAction a)
 {
   return (uint64_t)a;

src/types.h

@@ -72,8 +72,32 @@ enum class ProbeType
   hardware,
 };
 
+struct ProbeItem
+{
+  std::string name;
+  std::string abbr;
+  ProbeType type;
+};
+
+const ProbeItem PROBE_LIST[] =
+{
+  { "kprobe", "k", ProbeType::kprobe },
+  { "kretprobe", "kr", ProbeType::kretprobe },
+  { "uprobe", "u", ProbeType::uprobe },
+  { "uretprobe", "ur", ProbeType::uretprobe },
+  { "usdt", "U", ProbeType::usdt },
+  { "BEGIN", "BEGIN", ProbeType::uprobe },
+  { "END", "END", ProbeType::uprobe },
+  { "tracepoint", "t", ProbeType::tracepoint },
+  { "profile", "p", ProbeType::profile },
+  { "interval", "i", ProbeType::interval },
+  { "software", "s", ProbeType::software },
+  { "hardware", "h", ProbeType::hardware }
+};
+
 std::string typestr(Type t);
 ProbeType probetype(const std::string &type);
+std::string probetypeName(const std::string &type);
 
 class Probe
 {

tests/CMakeLists.txt

@@ -9,6 +9,7 @@ add_executable(bpftrace_test
   codegen.cpp
   main.cpp
   parser.cpp
+  probe.cpp
   semantic_analyser.cpp
   tracepoint_format_parser.cpp
   utils.cpp

tests/probe.cpp

+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "bpforc.h"
+#include "bpftrace.h"
+#include "clang_parser.h"
+#include "codegen_llvm.h"
+#include "driver.h"
+#include "fake_map.h"
+#include "semantic_analyser.h"
+
+namespace bpftrace
+{
+namespace test
+{
+namespace probe
+{
+
+using bpftrace::ast::AttachPoint;
+using bpftrace::ast::AttachPointList;
+using bpftrace::ast::Probe;
+
+void gen_bytecode(const std::string &input, std::stringstream &out)
+{
+	BPFtrace bpftrace;
+	Driver driver;
+	FakeMap::next_mapfd_ = 1;
+
+	ASSERT_EQ(driver.parse_str(input), 0);
+
+	ClangParser clang;
+	clang.parse(driver.root_, bpftrace.structs_);
+
+	ast::SemanticAnalyser semantics(driver.root_, bpftrace);
+	ASSERT_EQ(semantics.analyse(), 0);
+	ASSERT_EQ(semantics.create_maps(true), 0);
+
+	ast::CodegenLLVM codegen(driver.root_, bpftrace);
+	codegen.compile(DebugLevel::kDebug, out);
+}
+
+void compare_bytecode(const std::string &input1, const std::string &input2)
+{
+	std::stringstream expected_output1;
+	std::stringstream expected_output2;
+
+	gen_bytecode(input1, expected_output1);
+	gen_bytecode(input2, expected_output2);
+
+	EXPECT_EQ(expected_output1.str(), expected_output2.str());
+}
+
+TEST(probe, short_name)
+{
+	compare_bytecode("tracepoint:a:b { args }", "t:a:b { args }");
+	compare_bytecode("kprobe:f { pid }", "k:f { pid }");
+	compare_bytecode("kretprobe:f { pid }", "kr:f { pid }");
+	compare_bytecode("uprobe:path:f { 1 }", "u:path:f { 1 }");
+	compare_bytecode("profile:hz:997 { 1 }", "p:hz:997 { 1 }");
+	compare_bytecode("hardware:cache-references:1000000 { 1 }", "h:cache-references:1000000 { 1 }");
+	compare_bytecode("software:faults:1000 { 1 }", "s:faults:1000 { 1 }");
+	compare_bytecode("interval:s:1 { 1 }", "i:s:1 { 1 }");
+}
+
+} // namespace probe
+} // namespace test
+} // namespace bpftrace

tests/semantic_analyser.cpp

@@ -638,6 +638,19 @@ TEST(semantic_analyser, field_access_is_internal)
   EXPECT_EQ(true, var_assignment2->var->type.is_internal);
 }
 
+TEST(semantic_analyser, probe_short_name)
+{
+  test("t:a:b { args }", 0);
+  test("k:f { pid }", 0);
+  test("kr:f { pid }", 0);
+  test("u:path:f { 1 }", 0);
+  test("ur:path:f { 1 }", 0);
+  test("p:hz:997 { 1 }", 0);
+  test("h:cache-references:1000000 { 1 }", 0);
+  test("s:faults:1000 { 1 }", 0);
+  test("i:s:1 { 1 }", 0);
+}
+
 } // namespace semantic_analyser
 } // namespace test
 } // namespace bpftrace