1# Copyright 2023 Google LLC
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15from typing import Optional, List, Dict, Union, Tuple, Set
16from .ast import *
17
18
19def desugar_field_(field: Field, previous: Field, constraints: Dict[str, Constraint]) -> List[Field]:
20    """Inline group and constrained fields.
21    Constrained fields are transformed into fixed fields.
22    Group fields are inlined and recursively desugared."""
23
24    if isinstance(field, ScalarField) and field.id in constraints:
25        value = constraints[field.id].value
26        fixed = FixedField(kind='fixed_field', loc=field.loc, width=field.width, value=value)
27        fixed.parent = field.parent
28        return [fixed]
29
30    elif isinstance(field, PaddingField):
31        previous.padded_size = field.size
32        field.padded_field = previous
33        return [field]
34
35    elif isinstance(field, TypedefField) and field.id in constraints:
36        tag_id = constraints[field.id].tag_id
37        fixed = FixedField(kind='fixed_field', loc=field.loc, enum_id=field.type_id, tag_id=tag_id)
38        fixed.parent = field.parent
39        return [fixed]
40
41    elif isinstance(field, GroupField):
42        group = field.parent.file.group_scope[field.group_id]
43        constraints = dict([(c.id, c) for c in field.constraints])
44        fields = []
45        for f in group.fields:
46            fields.extend(desugar_field_(f, previous, constraints))
47            previous = f
48        return fields
49
50    else:
51        return [field]
52
53
54def desugar(file: File):
55    """Inline group fields.
56    Constrained fields are transformed into fixed fields.
57    Group declarations are removed from the file object.
58    **The original file object is modified inline.**"""
59
60    declarations = []
61    for d in file.declarations:
62        if isinstance(d, GroupDeclaration):
63            continue
64
65        if isinstance(d, (PacketDeclaration, StructDeclaration)):
66            fields = []
67            for f in d.fields:
68                fields.extend(desugar_field_(f, fields[-1] if len(fields) > 0 else None, {}))
69            d.fields = fields
70
71            # Add backlinks from optional links to their flag scalar
72            # field declaration.
73            fields_by_id = dict()
74            for f in d.fields:
75                if hasattr(f, 'id'):
76                    fields_by_id[f.id] = f
77                if f.cond:
78                    fields_by_id[f.cond.id].cond_for = f
79
80        declarations.append(d)
81
82    file.declarations = declarations
83    file.group_scope = {}
84
85
86def make_reserved_field(width: int) -> ReservedField:
87    """Create a reserved field of specified width."""
88    return ReservedField(kind='reserved_field', loc=None, width=width)
89
90
91def get_packet_field(packet: Union[PacketDeclaration, StructDeclaration], id: str) -> Optional[Field]:
92    """Return the field with selected identifier declared in the provided
93    packet or its ancestors."""
94    id = '_payload_' if id == 'payload' else id
95    for f in packet.fields:
96        if getattr(f, 'id', None) == id:
97            return f
98    if isinstance(packet, PacketDeclaration) and packet.parent_id:
99        parent = packet.file.packet_scope[packet.parent_id]
100        return get_packet_field(parent, id)
101    elif isinstance(packet, StructDeclaration) and packet.parent_id:
102        parent = packet.file.typedef_scope[packet.parent_id]
103        return get_packet_field(parent, id)
104    else:
105        return None
106
107
108def get_packet_fields(decl: Union[PacketDeclaration, StructDeclaration]) -> List[Field]:
109    """Return the list of fields declared in the selected packet and its parents.
110    Payload fields are removed from the parent declarations."""
111
112    fields = []
113    if decl.parent:
114        fields = [f for f in get_packet_fields(decl.parent) if not isinstance(f, (PayloadField, BodyField))]
115    return fields + decl.fields
116
117
118def get_packet_shift(packet: Union[PacketDeclaration, StructDeclaration]) -> int:
119    """Return the bit shift of the payload or body field in the parent packet.
120
121    When using packet derivation on bit fields, the body may be shifted.
122    The shift is handled statically in the implementation of child packets,
123    and the incomplete field is included in the body.
124    ```
125    packet Basic {
126        type: 1,
127        _body_
128    }
129    ```
130    """
131
132    # Traverse empty parents.
133    parent = packet.parent
134    while parent and len(parent.fields) == 1:
135        parent = parent.parent
136
137    if not parent:
138        return 0
139
140    shift = 0
141    for f in packet.parent.fields:
142        if isinstance(f, (BodyField, PayloadField)):
143            return 0 if (shift % 8) == 0 else shift
144        else:
145            # Fields that do not have a constant size are assumed to start
146            # on a byte boundary, and measure an integral number of bytes.
147            # Start the count over.
148            size = get_field_size(f)
149            shift = 0 if size is None else shift + size
150
151    # No payload or body in parent packet.
152    # Not raising an error, the generation will fail somewhere else.
153    return 0
154
155
156def get_packet_ancestor(
157        decl: Union[PacketDeclaration, StructDeclaration]) -> Union[PacketDeclaration, StructDeclaration]:
158    """Return the root ancestor of the selected packet or struct."""
159    if decl.parent_id is None:
160        return decl
161    else:
162        return get_packet_ancestor(decl.file.packet_scope[decl.parent_id])
163
164
165def get_derived_packets(
166    decl: Union[PacketDeclaration, StructDeclaration],
167    traverse: bool = True,
168) -> List[Tuple[List[Constraint], Union[PacketDeclaration, StructDeclaration]]]:
169    """Return the list of packets or structs that immediately derive from the
170    selected packet or struct, coupled with the field constraints.
171    Packet aliases (containing no field declarations other than a payload)
172    are traversed."""
173
174    children = []
175    for d in decl.file.declarations:
176        if type(d) is type(decl) and d.parent_id == decl.id:
177            if (len(d.fields) == 1 and isinstance(d.fields[0], (PayloadField, BodyField))) and traverse:
178                children.extend([(d.constraints + sub_constraints, sub_child)
179                                 for (sub_constraints, sub_child) in get_derived_packets(d)])
180            else:
181                children.append((d.constraints, d))
182    return children
183
184
185def get_all_packet_constraints(
186    decl: Union[PacketDeclaration, StructDeclaration]
187) -> List[Constraint]:
188    """Return the list of constraints defined in the selected packet and
189    its parent declarations."""
190
191    constraints = []
192    while decl.parent_id:
193        constraints.extend(decl.constraints)
194        decl = decl.parent
195    return constraints
196
197
198def get_field_size(field: Field, skip_payload: bool = False) -> Optional[int]:
199    """Determine the size of a field in bits, if possible.
200    If the field is dynamically sized (e.g. unsized array or payload field),
201    None is returned instead. If skip_payload is set, payload and body fields
202    are counted as having size 0 rather than a variable size."""
203
204    if field.cond:
205        return None
206
207    elif isinstance(field, (ScalarField, SizeField, CountField, ReservedField)):
208        return field.width
209
210    elif isinstance(field, FixedField):
211        return field.width or field.type.width
212
213    elif isinstance(field, PaddingField):
214        # Padding field width is added to the padded field size.
215        return 0
216
217    elif isinstance(field, ArrayField) and field.padded_size is not None:
218        return field.padded_size * 8
219
220    elif isinstance(field, ArrayField) and field.size is not None:
221        element_width = field.width or get_declaration_size(field.type)
222        return element_width * field.size if element_width is not None else None
223
224    elif isinstance(field, TypedefField):
225        return get_declaration_size(field.type)
226
227    elif isinstance(field, ChecksumField):
228        return 0
229
230    elif isinstance(field, (PayloadField, BodyField)) and skip_payload:
231        return 0
232
233    else:
234        return None
235
236
237def get_declaration_size(decl: Declaration, skip_payload: bool = False) -> Optional[int]:
238    """Determine the size of a declaration type in bits, if possible.
239    If the type is dynamically sized (e.g. contains an array or payload),
240    None is returned instead. If skip_payload is set, payload and body fields
241    are counted as having size 0 rather than a variable size."""
242
243    if isinstance(decl, (EnumDeclaration, CustomFieldDeclaration, ChecksumDeclaration)):
244        return decl.width
245
246    elif isinstance(decl, (PacketDeclaration, StructDeclaration)):
247        parent = decl.parent
248        packet_size = get_declaration_size(parent, skip_payload=True) if parent else 0
249        if packet_size is None:
250            return None
251        for f in decl.fields:
252            field_size = get_field_size(f, skip_payload=skip_payload)
253            if field_size is None:
254                return None
255            packet_size += field_size
256        return packet_size
257
258    else:
259        return None
260
261
262def get_array_field_size(field: ArrayField) -> Union[None, int, Field]:
263    """Return the array static size, size field, or count field.
264    If the array is unsized None is returned instead."""
265
266    if field.size is not None:
267        return field.size
268    for f in field.parent.fields:
269        if isinstance(f, (SizeField, CountField)) and f.field_id == field.id:
270            return f
271    return None
272
273
274def get_payload_field_size(field: Union[PayloadField, BodyField]) -> Optional[Field]:
275    """Return the payload or body size field.
276    If the payload is unsized None is returned instead."""
277
278    for f in field.parent.fields:
279        if isinstance(f, SizeField) and f.field_id == field.id:
280            return f
281    return None
282
283
284def get_array_element_size(field: ArrayField) -> Optional[int]:
285    """Return the array element size, if possible.
286    If the element size is not known at compile time,
287    None is returned instead."""
288
289    return field.width or get_declaration_size(field.type)
290
291
292def get_field_offset_from_start(field: Field) -> Optional[int]:
293    """Return the field bit offset from the start of the parent packet, if it
294    can be statically computed. If the offset is variable None is returned
295    instead."""
296    offset = 0
297    field_index = field.parent.fields.index(field)
298    for f in field.parent.fields[:field_index]:
299        size = get_field_size(f)
300        if size is None:
301            return None
302
303        offset += size
304    return offset
305
306
307def get_field_offset_from_end(field: Field) -> Optional[int]:
308    """Return the field bit offset from the end of the parent packet, if it
309    can be statically computed. If the offset is variable None is returned
310    instead. The selected field size is not counted towards the offset."""
311    offset = 0
312    field_index = field.parent.fields.index(field)
313    for f in field.parent.fields[field_index + 1:]:
314        size = get_field_size(f)
315        if size is None:
316            return None
317        offset += size
318    return offset
319
320
321def get_unconstrained_parent_fields(decl: Union[PacketDeclaration, StructDeclaration]) -> List[Field]:
322    """Return the list of fields from the parent declarations that have an identifier
323    but that do not have a value fixed by any of the parent constraints.
324    The fields are returned in order of declaration."""
325
326    def constraint_ids(constraints: List[Constraint]) -> Set[str]:
327        return set([c.id for c in constraints])
328
329    def aux(decl: Optional[Declaration], constraints: Set[str]) -> List[Field]:
330        if decl is None:
331            return []
332        fields = aux(decl.parent, constraints.union(constraint_ids(decl.constraints)))
333        for f in decl.fields:
334            if (isinstance(f, (ScalarField, ArrayField, TypedefField)) and not f.id in constraints):
335                fields.append(f)
336        return fields
337
338    return aux(decl.parent, constraint_ids(decl.constraints))
339
340
341def get_parent_constraints(decl: Union[PacketDeclaration, StructDeclaration]) -> List[Constraint]:
342    """Return the list of constraints from the current and parent declarations."""
343    parent_constraints = get_parent_constraints(decl.parent) if decl.parent else []
344    return parent_constraints + decl.constraints
345
346
347def is_bit_field(field: Field) -> bool:
348    """Identify fields that can have bit granularity.
349    These include: ScalarField, FixedField, TypedefField with enum type,
350    SizeField, and CountField."""
351
352    if field.cond:
353        return False
354
355    elif isinstance(field, (ScalarField, SizeField, CountField, FixedField, ReservedField)):
356        return True
357
358    elif isinstance(field, TypedefField) and isinstance(field.type, EnumDeclaration):
359        return True
360
361    else:
362        return False
363
364def is_open_enum(decl: EnumDeclaration) -> bool:
365    """Return true if the enum declaration is open, i.e. contains a tag
366    declaration of the form DEFAULT = .."""
367
368    return any(t.value is None for t in decl.tags)
369