Merge pull request #1385 from pchaigno/support-enums-map-visitor

Fix segfault with enumerations

src/cc/json_map_decl_visitor.cc

@@ -40,6 +40,8 @@ class BMapDeclVisitor : public clang::RecursiveASTVisitor<BMapDeclVisitor> {
   bool VisitTypedefType(const clang::TypedefType *T);
   bool VisitTagType(const clang::TagType *T);
   bool VisitPointerType(const clang::PointerType *T);
+  bool VisitEnumConstantDecl(clang::EnumConstantDecl *D);
+  bool VisitEnumDecl(clang::EnumDecl *D);
 
  private:
   clang::ASTContext &C;
@@ -55,6 +57,25 @@ bool BMapDeclVisitor::VisitFieldDecl(FieldDecl *D) {
   return true;
 }
 
+bool BMapDeclVisitor::VisitEnumConstantDecl(EnumConstantDecl *D) {
+  result_ += "\"";
+  result_ += D->getName();
+  result_ += "\",";
+  return false;
+}
+
+bool BMapDeclVisitor::VisitEnumDecl(EnumDecl *D) {
+  result_ += "[\"";
+  result_ += D->getName();
+  result_ += "\", [";
+  for (auto it = D->enumerator_begin(); it != D->enumerator_end(); ++it) {
+    TraverseDecl(*it);
+  }
+  result_.erase(result_.end() - 1);
+  result_ += "], \"enum\"]";
+  return false;
+}
+
 bool BMapDeclVisitor::TraverseRecordDecl(RecordDecl *D) {
   // skip children, handled in Visit...
   if (!WalkUpFromRecordDecl(D))

tests/python/test_clang.py

@@ -612,6 +612,18 @@ struct key_t {
         with self.assertRaises(Exception):
             b = BPF(text=text)
 
+    def test_enumerations(self):
+        text = """
+enum b {
+    CHOICE_A,
+};
+struct a {
+    enum b test;
+};
+BPF_HASH(drops, struct a);
+        """
+        b = BPF(text=text)
+
 
 if __name__ == "__main__":
     main()