Add _deref_to_depth in expr_pass

pythonbpf/expr_pass.py

@@ -129,10 +129,51 @@ def _handle_ctypes_call(
     return val
 
 
+def _get_base_type_and_depth(ir_type):
+    """Get the base type for pointer types."""
+    cur_type = ir_type
+    depth = 0
+    while isinstance(cur_type, ir.PointerType):
+        depth += 1
+        cur_type = cur_type.pointee
+    return cur_type, depth
+
+
+def _deref_to_depth(builder, val, target_depth):
+    """Dereference a pointer to a certain depth."""
+
+    cur_val = val
+    for _ in range(target_depth):
+        if not isinstance(val.type, ir.PointerType):
+            logger.error("Cannot dereference further, non-pointer type")
+            return None
+        cur_val = builder.load(cur_val)
+    return cur_val
+
+
+def _normalize_types(builder, lhs, rhs):
+    """Normalize types for comparison."""
+
+    if isinstance(lhs.type, ir.IntType) and isinstance(rhs.type, ir.IntType):
+        if lhs.type.width < rhs.type.width:
+            lhs = builder.sext(lhs, rhs.type)
+        else:
+            rhs = builder.sext(rhs, lhs.type)
+        return lhs, rhs
+
+    logger.error(f"Type mismatch: {lhs.type} vs {rhs.type}")
+    return None, None
+
+
 def _handle_comparator(builder, op, lhs, rhs):
     """Handle comparison operations."""
 
     # NOTE: For now assume same types
+    if lhs.type != rhs.type:
+        lhs, rhs = _normalize_types(builder, lhs, rhs)
+
+    if lhs is None or rhs is None:
+        return None
 
     comparison_ops = {
         ast.Eq: "==",