Add support for unary op 'not' in eval_expr, move not test to passing

Author: r41k0u

pythonbpf/expr_pass.py

@@ -22,7 +22,6 @@ def _handle_name_expr(expr: ast.Name, local_sym_tab: Dict, builder: ir.IRBuilder
 
 def _handle_constant_expr(expr: ast.Constant):
     """Handle ast.Constant expressions."""
-    logger.info("We the best")
     if isinstance(expr.value, int) or isinstance(expr.value, bool):
         return ir.Constant(ir.IntType(64), int(expr.value)), ir.IntType(64)
     else:
@@ -191,6 +190,43 @@ def _handle_compare(
     return _handle_comparator(builder, cond.ops[0], lhs, rhs)
 
 
+def convert_to_bool(builder, val):
+    if val.type == ir.IntType(1):
+        return val
+    if isinstance(val.type, ir.PointerType):
+        zero = ir.Constant(val.type, None)
+    else:
+        zero = ir.Constant(val.type, 0)
+    return builder.icmp_signed("!=", val, zero)
+
+
+def _handle_unary_op(
+    func,
+    module,
+    builder,
+    expr: ast.UnaryOp,
+    local_sym_tab,
+    map_sym_tab,
+    structs_sym_tab=None,
+):
+    """Handle ast.UnaryOp expressions."""
+    if not isinstance(expr.op, ast.Not):
+        logger.error("Only 'not' unary operator is supported")
+        return None
+
+    operand = eval_expr(
+        func, module, builder, expr.operand, local_sym_tab, map_sym_tab, structs_sym_tab
+    )
+    if operand is None:
+        logger.error("Failed to evaluate operand for unary operation")
+        return None
+
+    operand_val, operand_type = operand
+    true_const = ir.Constant(ir.IntType(1), 1)
+    result = builder.xor(convert_to_bool(builder, operand_val), true_const)
+    return result, ir.IntType(1)
+
+
 def eval_expr(
     func,
     module,
@@ -275,6 +311,10 @@ def eval_expr(
         return _handle_compare(
             func, module, builder, expr, local_sym_tab, map_sym_tab, structs_sym_tab
         )
+    elif isinstance(expr, ast.UnaryOp):
+        return _handle_unary_op(
+            func, module, builder, expr, local_sym_tab, map_sym_tab, structs_sym_tab
+        )
     logger.info("Unsupported expression evaluation")
     return None
 

pythonbpf/functions/functions_pass.py

@@ -7,7 +7,7 @@
 from pythonbpf.helper import HelperHandlerRegistry, handle_helper_call
 from pythonbpf.type_deducer import ctypes_to_ir
 from pythonbpf.binary_ops import handle_binary_op
-from pythonbpf.expr_pass import eval_expr, handle_expr
+from pythonbpf.expr_pass import eval_expr, handle_expr, convert_to_bool
 
 from .return_utils import _handle_none_return, _handle_xdp_return, _is_xdp_name
 
@@ -240,20 +240,10 @@ def handle_assign(
         logger.info("Unsupported assignment value type")
 
 
-def _convert_to_bool(builder, val):
-    if val.type == ir.IntType(1):
-        return val
-    if isinstance(val.type, ir.PointerType):
-        zero = ir.Constant(val.type, None)
-    else:
-        zero = ir.Constant(val.type, 0)
-    return builder.icmp_signed("!=", val, zero)
-
-
 def handle_cond(func, module, builder, cond, local_sym_tab, map_sym_tab):
     if True:
         val = eval_expr(func, module, builder, cond, local_sym_tab, map_sym_tab)[0]
-        return _convert_to_bool(builder, val)
+        return convert_to_bool(builder, val)
     if isinstance(cond, ast.Constant):
         if isinstance(cond.value, bool) or isinstance(cond.value, int):
             return ir.Constant(ir.IntType(1), int(cond.value))