xref: /llvm-project/lldb/test/API/lang/rust/enum-structs/RustEnumValue.py (revision 6bf923d5c3daf5d66e0acf53d037a12ceff4a275) 
1"""Helper library to traverse data emitted for Rust enums """
2from lldbsuite.test.lldbtest import *
3
4DISCRIMINANT_MEMBER_NAME = "$discr$"
5VALUE_MEMBER_NAME = "value"
6
7
8class RustEnumValue:
9    def __init__(self, value: lldb.SBValue):
10        self.value = value
11
12    def getAllVariantTypes(self):
13        result = []
14        for i in range(self._inner().GetNumChildren()):
15            result.append(self.getVariantByIndex(i).GetDisplayTypeName())
16        return result
17
18    def _inner(self) -> lldb.SBValue:
19        return self.value.GetChildAtIndex(0)
20
21    def getVariantByIndex(self, index):
22        return (
23            self._inner()
24            .GetChildAtIndex(index)
25            .GetChildMemberWithName(VALUE_MEMBER_NAME)
26        )
27
28    @staticmethod
29    def _getDiscriminantValueAsUnsigned(discr_sbvalue: lldb.SBValue):
30        byte_size = discr_sbvalue.GetType().GetByteSize()
31        error = lldb.SBError()
32
33        # when discriminant is u16 Clang emits 'unsigned char'
34        # and LLDB seems to treat it as character type disalowing to call GetValueAsUnsigned
35        if byte_size == 1:
36            return discr_sbvalue.GetData().GetUnsignedInt8(error, 0)
37        elif byte_size == 2:
38            return discr_sbvalue.GetData().GetUnsignedInt16(error, 0)
39        elif byte_size == 4:
40            return discr_sbvalue.GetData().GetUnsignedInt32(error, 0)
41        elif byte_size == 8:
42            return discr_sbvalue.GetData().GetUnsignedInt64(error, 0)
43        else:
44            return discr_sbvalue.GetValueAsUnsigned()
45
46    def getCurrentVariantIndex(self):
47        default_index = 0
48        for i in range(self._inner().GetNumChildren()):
49            variant: lldb.SBValue = self._inner().GetChildAtIndex(i)
50            discr = variant.GetChildMemberWithName(DISCRIMINANT_MEMBER_NAME)
51            if discr.IsValid():
52                discr_unsigned_value = RustEnumValue._getDiscriminantValueAsUnsigned(
53                    discr
54                )
55                if variant.GetName() == f"$variant${discr_unsigned_value}":
56                    return discr_unsigned_value
57            else:
58                default_index = i
59        return default_index
60
61    def getFields(self):
62        result = []
63        for i in range(self._inner().GetNumChildren()):
64            type: lldb.SBType = self._inner().GetType()
65            result.append(type.GetFieldAtIndex(i).GetName())
66        return result
67
68    def getCurrentValue(self) -> lldb.SBValue:
69        return self.getVariantByIndex(self.getCurrentVariantIndex())
70