Skip to content

Commit 44f39da

Browse files
jhonabreuljhonabreul
committed
Refactor enums comparison operators for performance improvements
1 parent 1d09c98 commit 44f39da

File tree

4 files changed

+276
-460
lines changed

4 files changed

+276
-460
lines changed

‎src/runtime/ClassManager.cs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,11 @@ internal static ClassBase CreateClass(Type type)
220220
impl = new LookUpObject(type);
221221
}
222222

223+
else if (type.IsEnum)
224+
{
225+
impl = new EnumObject(type);
226+
}
227+
223228
else
224229
{
225230
impl = new ClassObject(type);

‎src/runtime/Types/ClassBase.cs‎

Lines changed: 48 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -156,42 +156,7 @@ public static NewReference tp_richcompare(BorrowedReference ob, BorrowedReferenc
156156
try
157157
{
158158
int cmp = co1Comp.CompareTo(co2Inst);
159-
160-
BorrowedReference pyCmp;
161-
if (cmp < 0)
162-
{
163-
if (op == Runtime.Py_LT || op == Runtime.Py_LE)
164-
{
165-
pyCmp = Runtime.PyTrue;
166-
}
167-
else
168-
{
169-
pyCmp = Runtime.PyFalse;
170-
}
171-
}
172-
else if (cmp == 0)
173-
{
174-
if (op == Runtime.Py_LE || op == Runtime.Py_GE)
175-
{
176-
pyCmp = Runtime.PyTrue;
177-
}
178-
else
179-
{
180-
pyCmp = Runtime.PyFalse;
181-
}
182-
}
183-
else
184-
{
185-
if (op == Runtime.Py_GE || op == Runtime.Py_GT)
186-
{
187-
pyCmp = Runtime.PyTrue;
188-
}
189-
else
190-
{
191-
pyCmp = Runtime.PyFalse;
192-
}
193-
}
194-
return new NewReference(pyCmp);
159+
return new NewReference(GetComparisonResult(op, cmp));
195160
}
196161
catch (ArgumentException e)
197162
{
@@ -202,7 +167,53 @@ public static NewReference tp_richcompare(BorrowedReference ob, BorrowedReferenc
202167
}
203168
}
204169

205-
private static bool TryGetSecondCompareOperandInstance(BorrowedReference left, BorrowedReference right, out CLRObject co1, out object co2Inst)
170+
/// <summary>
171+
/// Get the result of a comparison operation based on the operator and the comparison result.
172+
/// </summary>
173+
/// <remarks>
174+
/// This method is used to determine the result of a comparison operation, excluding equality and inequality.
175+
/// </remarks>
176+
protected static BorrowedReference GetComparisonResult(int op, int comparisonResult)
177+
{
178+
BorrowedReference pyCmp;
179+
if (comparisonResult < 0)
180+
{
181+
if (op == Runtime.Py_LT || op == Runtime.Py_LE)
182+
{
183+
pyCmp = Runtime.PyTrue;
184+
}
185+
else
186+
{
187+
pyCmp = Runtime.PyFalse;
188+
}
189+
}
190+
else if (comparisonResult == 0)
191+
{
192+
if (op == Runtime.Py_LE || op == Runtime.Py_GE)
193+
{
194+
pyCmp = Runtime.PyTrue;
195+
}
196+
else
197+
{
198+
pyCmp = Runtime.PyFalse;
199+
}
200+
}
201+
else
202+
{
203+
if (op == Runtime.Py_GE || op == Runtime.Py_GT)
204+
{
205+
pyCmp = Runtime.PyTrue;
206+
}
207+
else
208+
{
209+
pyCmp = Runtime.PyFalse;
210+
}
211+
}
212+
213+
return pyCmp;
214+
}
215+
216+
protected static bool TryGetSecondCompareOperandInstance(BorrowedReference left, BorrowedReference right, out CLRObject co1, out object co2Inst)
206217
{
207218
co2Inst = null;
208219

‎src/runtime/Types/EnumObject.cs‎

Lines changed: 223 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,223 @@
1+
using System;
2+
using System.Runtime.CompilerServices;
3+
4+
namespace Python.Runtime
5+
{
6+
/// <summary>
7+
/// Managed class that provides the implementation for reflected enum types.
8+
/// </summary>
9+
[Serializable]
10+
internal class EnumObject : ClassObject
11+
{
12+
private bool _isUnsigned;
13+
14+
internal EnumObject(Type type) : base(type)
15+
{
16+
if (!type.IsEnum)
17+
{
18+
throw new ArgumentException($"{nameof(EnumObject)} can only handle Enum types. Received {type.Name}", nameof(type));
19+
}
20+
21+
_isUnsigned = type.GetEnumUnderlyingType() == typeof(UInt64);
22+
}
23+
24+
/// <summary>
25+
/// Standard comparison implementation for instances of enum types.
26+
/// </summary>
27+
public new static NewReference tp_richcompare(BorrowedReference ob, BorrowedReference other, int op)
28+
{
29+
object rightInstance;
30+
CLRObject leftClrObject;
31+
int comparisonResult;
32+
var leftType = Runtime.PyObject_TYPE(ob);
33+
var leftClass = (EnumObject)GetManagedObject(leftType)!;
34+
35+
switch (op)
36+
{
37+
case Runtime.Py_EQ:
38+
case Runtime.Py_NE:
39+
var pytrue = Runtime.PyTrue;
40+
var pyfalse = Runtime.PyFalse;
41+
42+
// swap true and false for NE
43+
if (op != Runtime.Py_EQ)
44+
{
45+
pytrue = Runtime.PyFalse;
46+
pyfalse = Runtime.PyTrue;
47+
}
48+
49+
if (ob == other)
50+
{
51+
return new NewReference(pytrue);
52+
}
53+
54+
if (!TryGetSecondCompareOperandInstance(ob, other, out leftClrObject, out rightInstance))
55+
{
56+
return new NewReference(pyfalse);
57+
}
58+
59+
if (rightInstance != null &&
60+
TryCompare(leftClrObject.inst as Enum, rightInstance, leftClass._isUnsigned, out comparisonResult) &&
61+
comparisonResult == 0)
62+
{
63+
return new NewReference(pytrue);
64+
}
65+
else
66+
{
67+
return new NewReference(pyfalse);
68+
}
69+
70+
case Runtime.Py_LT:
71+
case Runtime.Py_LE:
72+
case Runtime.Py_GT:
73+
case Runtime.Py_GE:
74+
if (!TryGetSecondCompareOperandInstance(ob, other, out leftClrObject, out rightInstance))
75+
{
76+
return Exceptions.RaiseTypeError("Cannot get managed object");
77+
}
78+
79+
if (rightInstance == null)
80+
{
81+
return Exceptions.RaiseTypeError($"Cannot compare {leftClrObject.inst.GetType()} to None");
82+
}
83+
84+
try
85+
{
86+
if (!TryCompare(leftClrObject.inst as Enum, rightInstance, leftClass._isUnsigned, out comparisonResult))
87+
{
88+
return Exceptions.RaiseTypeError($"Cannot compare {leftClrObject.inst.GetType()} with {rightInstance.GetType()}");
89+
}
90+
91+
return new NewReference(GetComparisonResult(op, comparisonResult));
92+
}
93+
catch (ArgumentException e)
94+
{
95+
return Exceptions.RaiseTypeError(e.Message);
96+
}
97+
98+
default:
99+
return new NewReference(Runtime.PyNotImplemented);
100+
}
101+
}
102+
103+
/// <summary>
104+
/// Tries comparing the give enum to the right operand by converting it to the appropriate type if possible
105+
/// </summary>
106+
/// <returns>True if the right operand was converted to a supported type and the comparison was performed successfully</returns>
107+
private static bool TryCompare(Enum left, object right, bool isUnsigned, out int result)
108+
{
109+
result = int.MinValue;
110+
var conversionSuccessful = true;
111+
// Same enum comparison:
112+
if (left.GetType() == right.GetType())
113+
{
114+
result = left.CompareTo(right);
115+
}
116+
// Comparison with other enum types
117+
else if (right.GetType().IsEnum)
118+
{
119+
result = Compare(left, right as Enum, isUnsigned);
120+
}
121+
else if (right is double rightDouble)
122+
{
123+
result = Compare(left, rightDouble, isUnsigned);
124+
}
125+
else if (right is long rightLong)
126+
{
127+
result = Compare(left, rightLong, isUnsigned);
128+
}
129+
else if (right is ulong rightULong)
130+
{
131+
result = Compare(left, rightULong, isUnsigned);
132+
}
133+
else if (right is int rightInt)
134+
{
135+
result = Compare(left, rightInt, isUnsigned);
136+
}
137+
else if (right is uint rightUInt)
138+
{
139+
result = Compare(left, rightUInt, isUnsigned);
140+
}
141+
else if (right is string rightString)
142+
{
143+
result = left.ToString().CompareTo(rightString);
144+
}
145+
else
146+
{
147+
conversionSuccessful = false;
148+
}
149+
150+
return conversionSuccessful;
151+
}
152+
153+
#region Comparison against integers
154+
155+
[MethodImpl(MethodImplOptions.AggressiveInlining)]
156+
private static int Compare(long a, ulong b)
157+
{
158+
if (a < 0) return -1;
159+
return ((ulong)a).CompareTo(b);
160+
}
161+
162+
[MethodImpl(MethodImplOptions.AggressiveInlining)]
163+
private static int Compare(Enum a, long b, bool isUnsigned)
164+
{
165+
166+
if (isUnsigned)
167+
{
168+
return -Compare(b, Convert.ToUInt64(a));
169+
}
170+
return Convert.ToInt64(a).CompareTo(b);
171+
}
172+
173+
[MethodImpl(MethodImplOptions.AggressiveInlining)]
174+
private static int Compare(Enum a, ulong b, bool inUnsigned)
175+
{
176+
if (inUnsigned)
177+
{
178+
return Convert.ToUInt64(a).CompareTo(b);
179+
}
180+
return Compare(Convert.ToInt64(a), b);
181+
}
182+
183+
#endregion
184+
185+
#region Comparison against doubles
186+
187+
[MethodImpl(MethodImplOptions.AggressiveInlining)]
188+
private static int Compare(Enum a, double b, bool isUnsigned)
189+
{
190+
if (isUnsigned)
191+
{
192+
var uIntA = Convert.ToUInt64(a);
193+
if (uIntA < b) return -1;
194+
if (uIntA > b) return 1;
195+
return 0;
196+
}
197+
198+
var intA = Convert.ToInt64(a);
199+
if (intA < b) return -1;
200+
if (intA > b) return 1;
201+
return 0;
202+
}
203+
204+
#endregion
205+
206+
#region Comparison against other enum types
207+
208+
/// <summary>
209+
/// We support comparing enums of different types by comparing their underlying values.
210+
/// </summary>
211+
[MethodImpl(MethodImplOptions.AggressiveInlining)]
212+
private static int Compare(Enum a, Enum b, bool isUnsigned)
213+
{
214+
if (b.GetType().GetEnumUnderlyingType() == typeof(UInt64))
215+
{
216+
return Compare(a, Convert.ToUInt64(b), isUnsigned);
217+
}
218+
return Compare(a, Convert.ToInt64(b), isUnsigned);
219+
}
220+
221+
#endregion
222+
}
223+
}

0 commit comments

Comments
 (0)