1#------------------------------------------------------------------------------
2# pycparser: ast_transforms.py
3#
4# Some utilities used by the parser to create a friendlier AST.
5#
6# Eli Bendersky [https://eli.thegreenplace.net/]
7# License: BSD
8#------------------------------------------------------------------------------
9
10from . import c_ast
11
12
13def fix_switch_cases(switch_node):
14    """ The 'case' statements in a 'switch' come out of parsing with one
15        child node, so subsequent statements are just tucked to the parent
16        Compound. Additionally, consecutive (fall-through) case statements
17        come out messy. This is a peculiarity of the C grammar. The following:
18
19            switch (myvar) {
20                case 10:
21                    k = 10;
22                    p = k + 1;
23                    return 10;
24                case 20:
25                case 30:
26                    return 20;
27                default:
28                    break;
29            }
30
31        Creates this tree (pseudo-dump):
32
33            Switch
34                ID: myvar
35                Compound:
36                    Case 10:
37                        k = 10
38                    p = k + 1
39                    return 10
40                    Case 20:
41                        Case 30:
42                            return 20
43                    Default:
44                        break
45
46        The goal of this transform is to fix this mess, turning it into the
47        following:
48
49            Switch
50                ID: myvar
51                Compound:
52                    Case 10:
53                        k = 10
54                        p = k + 1
55                        return 10
56                    Case 20:
57                    Case 30:
58                        return 20
59                    Default:
60                        break
61
62        A fixed AST node is returned. The argument may be modified.
63    """
64    assert isinstance(switch_node, c_ast.Switch)
65    if not isinstance(switch_node.stmt, c_ast.Compound):
66        return switch_node
67
68    # The new Compound child for the Switch, which will collect children in the
69    # correct order
70    new_compound = c_ast.Compound([], switch_node.stmt.coord)
71
72    # The last Case/Default node
73    last_case = None
74
75    # Goes over the children of the Compound below the Switch, adding them
76    # either directly below new_compound or below the last Case as appropriate
77    # (for `switch(cond) {}`, block_items would have been None)
78    for child in (switch_node.stmt.block_items or []):
79        if isinstance(child, (c_ast.Case, c_ast.Default)):
80            # If it's a Case/Default:
81            # 1. Add it to the Compound and mark as "last case"
82            # 2. If its immediate child is also a Case or Default, promote it
83            #    to a sibling.
84            new_compound.block_items.append(child)
85            _extract_nested_case(child, new_compound.block_items)
86            last_case = new_compound.block_items[-1]
87        else:
88            # Other statements are added as children to the last case, if it
89            # exists.
90            if last_case is None:
91                new_compound.block_items.append(child)
92            else:
93                last_case.stmts.append(child)
94
95    switch_node.stmt = new_compound
96    return switch_node
97
98
99def _extract_nested_case(case_node, stmts_list):
100    """ Recursively extract consecutive Case statements that are made nested
101        by the parser and add them to the stmts_list.
102    """
103    if isinstance(case_node.stmts[0], (c_ast.Case, c_ast.Default)):
104        stmts_list.append(case_node.stmts.pop())
105        _extract_nested_case(stmts_list[-1], stmts_list)
106
107