Fix #235: Pragma displacing real statements (#236)

* Fix #235: Pragma displacing real statements

2 files changed, 111 insertions, 10 deletions

diff --git a/pycparser/c_parser.py b/pycparser/c_parser.py
index f01f67f..47d958f 100644
--- a/pycparser/c_parser.py
+++ b/pycparser/c_parser.py
@@ -616,6 +616,59 @@ class CParser(PLYParser):
         """
         p[0] = p[1]
 
+    # A pragma is generally considered a decorator rather than an actual statement.
+    # Still, for the purposes of analyzing an abstract syntax tree of C code,
+    # pragma's should not be ignored and were previously treated as a statement.
+    # This presents a problem for constructs that take a statement such as labeled_statements,
+    # selection_statements, and iteration_statements, causing a misleading structure
+    # in the AST. For example, consider the following C code.
+    #
+    #   for (int i = 0; i < 3; i++)
+    #       #pragma omp critical
+    #       sum += 1;
+    #
+    # This code will compile and execute "sum += 1;" as the body of the for loop.
+    # Previous implementations of PyCParser would render the AST for this
+    # block of code as follows:
+    #
+    #   For:
+    #     DeclList:
+    #       Decl: i, [], [], []
+    #         TypeDecl: i, []
+    #           IdentifierType: ['int']
+    #         Constant: int, 0
+    #     BinaryOp: <
+    #       ID: i
+    #       Constant: int, 3
+    #     UnaryOp: p++
+    #       ID: i
+    #     Pragma: omp critical
+    #   Assignment: +=
+    #     ID: sum
+    #     Constant: int, 1
+    #
+    # This AST misleadingly takes the Pragma as the body of the loop and the
+    # assignment then becomes a sibling of the loop.
+    #
+    # To solve edge cases like these, the pragmacomp_or_statement rule groups
+    # a pragma and its following statement (which would otherwise be orphaned)
+    # using a compound block, effectively turning the above code into:
+    #
+    #   for (int i = 0; i < 3; i++) {
+    #       #pragma omp critical
+    #       sum += 1;
+    #   }
+    def p_pragmacomp_or_statement(self, p):
+        """ pragmacomp_or_statement     : pppragma_directive statement
+                                        | statement
+        """
+        if isinstance(p[1], c_ast.Pragma) and len(p) == 3:
+            p[0] = c_ast.Compound(
+                block_items=[p[1], p[2]],
+                coord=self._token_coord(p, 1))
+        else:
+            p[0] = p[1]
+
     # In C, declarations can come several in a line:
     #   int x, *px, romulo = 5;
     #
@@ -1410,44 +1463,44 @@ class CParser(PLYParser):
             coord=self._token_coord(p, 1))
 
     def p_labeled_statement_1(self, p):
-        """ labeled_statement : ID COLON statement """
+        """ labeled_statement : ID COLON pragmacomp_or_statement """
         p[0] = c_ast.Label(p[1], p[3], self._token_coord(p, 1))
 
     def p_labeled_statement_2(self, p):
-        """ labeled_statement : CASE constant_expression COLON statement """
+        """ labeled_statement : CASE constant_expression COLON pragmacomp_or_statement """
         p[0] = c_ast.Case(p[2], [p[4]], self._token_coord(p, 1))
 
     def p_labeled_statement_3(self, p):
-        """ labeled_statement : DEFAULT COLON statement """
+        """ labeled_statement : DEFAULT COLON pragmacomp_or_statement """
         p[0] = c_ast.Default([p[3]], self._token_coord(p, 1))
 
     def p_selection_statement_1(self, p):
-        """ selection_statement : IF LPAREN expression RPAREN statement """
+        """ selection_statement : IF LPAREN expression RPAREN pragmacomp_or_statement """
         p[0] = c_ast.If(p[3], p[5], None, self._token_coord(p, 1))
 
     def p_selection_statement_2(self, p):
-        """ selection_statement : IF LPAREN expression RPAREN statement ELSE statement """
+        """ selection_statement : IF LPAREN expression RPAREN statement ELSE pragmacomp_or_statement """
         p[0] = c_ast.If(p[3], p[5], p[7], self._token_coord(p, 1))
 
     def p_selection_statement_3(self, p):
-        """ selection_statement : SWITCH LPAREN expression RPAREN statement """
+        """ selection_statement : SWITCH LPAREN expression RPAREN pragmacomp_or_statement """
         p[0] = fix_switch_cases(
                 c_ast.Switch(p[3], p[5], self._token_coord(p, 1)))
 
     def p_iteration_statement_1(self, p):
-        """ iteration_statement : WHILE LPAREN expression RPAREN statement """
+        """ iteration_statement : WHILE LPAREN expression RPAREN pragmacomp_or_statement """
         p[0] = c_ast.While(p[3], p[5], self._token_coord(p, 1))
 
     def p_iteration_statement_2(self, p):
-        """ iteration_statement : DO statement WHILE LPAREN expression RPAREN SEMI """
+        """ iteration_statement : DO pragmacomp_or_statement WHILE LPAREN expression RPAREN SEMI """
         p[0] = c_ast.DoWhile(p[5], p[2], self._token_coord(p, 1))
 
     def p_iteration_statement_3(self, p):
-        """ iteration_statement : FOR LPAREN expression_opt SEMI expression_opt SEMI expression_opt RPAREN statement """
+        """ iteration_statement : FOR LPAREN expression_opt SEMI expression_opt SEMI expression_opt RPAREN pragmacomp_or_statement """
         p[0] = c_ast.For(p[3], p[5], p[7], p[9], self._token_coord(p, 1))
 
     def p_iteration_statement_4(self, p):
-        """ iteration_statement : FOR LPAREN declaration expression_opt SEMI expression_opt RPAREN statement """
+        """ iteration_statement : FOR LPAREN declaration expression_opt SEMI expression_opt RPAREN pragmacomp_or_statement """
         p[0] = c_ast.For(c_ast.DeclList(p[3], self._token_coord(p, 1)),
                          p[4], p[6], p[8], self._token_coord(p, 1))
 
diff --git a/tests/test_c_parser.py b/tests/test_c_parser.py
index ab6143f..3b336bf 100755
--- a/tests/test_c_parser.py
+++ b/tests/test_c_parser.py
@@ -1369,6 +1369,54 @@ class TestCParser_fundamentals(TestCParser_base):
         self.assertEqual(s1_ast.ext[2].type.type.decls[0].string, 'baz')
         self.assertEqual(s1_ast.ext[2].type.type.decls[0].coord.line, 9)
 
+    def test_pragmacomp_or_statement(self):
+        s1 = r'''
+            void main() {
+                int sum = 0;
+                for (int i; i < 3; i++)
+                    #pragma omp critical
+                    sum += 1;
+
+                while(sum < 10)
+                    #pragma omp critical
+                    sum += 1;
+
+                mylabel:
+                    #pragma foo
+                    sum += 10;
+
+                if (sum > 10)
+                    #pragma bar
+                    sum = 10;
+
+                switch (sum)
+                case 10:
+                    #pragma foo
+                    sum = 20;
+            }
+        '''
+        s1_ast = self.parse(s1)
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[1], For))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[1].stmt, Compound))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[1].stmt.block_items[0], Pragma))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[1].stmt.block_items[1], Assignment))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[2], While))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[2].stmt, Compound))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[2].stmt.block_items[0], Pragma))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[2].stmt.block_items[1], Assignment))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[3], Label))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[3].stmt, Compound))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[3].stmt.block_items[0], Pragma))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[3].stmt.block_items[1], Assignment))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[4], If))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[4].iftrue, Compound))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[4].iftrue.block_items[0], Pragma))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[4].iftrue.block_items[1], Assignment))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[5], Switch))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[5].stmt.stmts[0], Compound))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[5].stmt.stmts[0].block_items[0], Pragma))
+        self.assertTrue(isinstance(s1_ast.ext[0].body.block_items[5].stmt.stmts[0].block_items[1], Assignment))
+
 
 class TestCParser_whole_code(TestCParser_base):
     """ Testing of parsing whole chunks of code.