feat: add short-circuit for || and &&

swananan · swananan · commit 23543bb834d4 · 2025-10-06T21:30:01.000+08:00
diff --git a/docs/scripting.md b/docs/scripting.md
@@ -236,7 +236,9 @@ let complex = (x + y) / (a - b);
 
 - `&&` (logical AND), `||` (logical OR)
 - Operands are treated as booleans with "non-zero is true" semantics
-- Current implementation evaluates both sides (no short-circuit yet)
+- `||` and `&&` use short-circuit evaluation
+  - `||`: if LHS is true, RHS is not evaluated
+  - `&&`: if LHS is false, RHS is not evaluated
 
 Boolean values
 
diff --git a/docs/zh/scripting.md b/docs/zh/scripting.md
@@ -288,7 +288,9 @@ let complex = (x + y) / (a - b);
 
 - `&&`（逻辑与）、`||`（逻辑或）
 - 操作数按“非零为真”处理
-- 当前实现为“非短路”：左右两侧都会被求值
+- `||`、`&&` 均采用短路求值：
+  - `||`：左侧为真则不再计算右侧
+  - `&&`：左侧为假则不再计算右侧
 
 示例
 
diff --git a/ghostscope-compiler/src/ebpf/expression.rs b/ghostscope-compiler/src/ebpf/expression.rs
@@ -78,6 +78,161 @@ impl<'ctx> EbpfContext<'ctx> {
             }
             Expr::SpecialVar(name) => self.handle_special_variable(name),
             Expr::BinaryOp { left, op, right } => {
+                // Implement short-circuit for logical OR (||) and logical AND (&&)
+                if matches!(op, BinaryOp::LogicalOr) {
+                    // Evaluate LHS to boolean (non-zero => true)
+                    let lhs_val = self.compile_expr(left)?;
+                    let lhs_int = match lhs_val {
+                        BasicValueEnum::IntValue(iv) => iv,
+                        _ => {
+                            return Err(CodeGenError::TypeError(
+                                "Logical OR requires integer operands".to_string(),
+                            ))
+                        }
+                    };
+                    let lhs_zero = lhs_int.get_type().const_zero();
+                    let lhs_bool = self
+                        .builder
+                        .build_int_compare(
+                            inkwell::IntPredicate::NE,
+                            lhs_int,
+                            lhs_zero,
+                            "lor_lhs_nz",
+                        )
+                        .map_err(|e| CodeGenError::Builder(e.to_string()))?;
+
+                    // Prepare control flow blocks
+                    let curr_block = self.builder.get_insert_block().ok_or_else(|| {
+                        CodeGenError::LLVMError("No current basic block".to_string())
+                    })?;
+                    let func = curr_block
+                        .get_parent()
+                        .ok_or_else(|| CodeGenError::LLVMError("No parent function".to_string()))?;
+                    let rhs_block = self.context.append_basic_block(func, "lor_rhs");
+                    let merge_block = self.context.append_basic_block(func, "lor_merge");
+
+                    // If lhs is true, jump directly to merge (short-circuit)
+                    self.builder
+                        .build_conditional_branch(lhs_bool, merge_block, rhs_block)
+                        .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
+
+                    // RHS path: compute boolean only if needed
+                    self.builder.position_at_end(rhs_block);
+                    let rhs_val = self.compile_expr(right)?;
+                    let rhs_int = match rhs_val {
+                        BasicValueEnum::IntValue(iv) => iv,
+                        _ => {
+                            return Err(CodeGenError::TypeError(
+                                "Logical OR requires integer operands".to_string(),
+                            ))
+                        }
+                    };
+                    let rhs_zero = rhs_int.get_type().const_zero();
+                    let rhs_bool = self
+                        .builder
+                        .build_int_compare(
+                            inkwell::IntPredicate::NE,
+                            rhs_int,
+                            rhs_zero,
+                            "lor_rhs_nz",
+                        )
+                        .map_err(|e| CodeGenError::Builder(e.to_string()))?;
+                    // Capture the actual block where RHS computation ended
+                    let rhs_end_block = self.builder.get_insert_block().ok_or_else(|| {
+                        CodeGenError::LLVMError("No current basic block after RHS".to_string())
+                    })?;
+                    self.builder
+                        .build_unconditional_branch(merge_block)
+                        .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
+
+                    // Merge: phi of i1: true from LHS-true, RHS bool from rhs_block
+                    self.builder.position_at_end(merge_block);
+                    let i1 = self.context.bool_type();
+                    let phi = self
+                        .builder
+                        .build_phi(i1, "lor_phi")
+                        .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
+                    let one = i1.const_int(1, false);
+                    phi.add_incoming(&[(&one, curr_block), (&rhs_bool, rhs_end_block)]);
+                    return Ok(phi.as_basic_value());
+                } else if matches!(op, BinaryOp::LogicalAnd) {
+                    // Evaluate LHS to boolean (non-zero => true)
+                    let lhs_val = self.compile_expr(left)?;
+                    let lhs_int = match lhs_val {
+                        BasicValueEnum::IntValue(iv) => iv,
+                        _ => {
+                            return Err(CodeGenError::TypeError(
+                                "Logical AND requires integer operands".to_string(),
+                            ))
+                        }
+                    };
+                    let lhs_zero = lhs_int.get_type().const_zero();
+                    let lhs_bool = self
+                        .builder
+                        .build_int_compare(
+                            inkwell::IntPredicate::NE,
+                            lhs_int,
+                            lhs_zero,
+                            "land_lhs_nz",
+                        )
+                        .map_err(|e| CodeGenError::Builder(e.to_string()))?;
+
+                    // Prepare control flow: if lhs is true, evaluate rhs; else short-circuit to false
+                    let curr_block = self.builder.get_insert_block().ok_or_else(|| {
+                        CodeGenError::LLVMError("No current basic block".to_string())
+                    })?;
+                    let func = curr_block
+                        .get_parent()
+                        .ok_or_else(|| CodeGenError::LLVMError("No parent function".to_string()))?;
+                    let rhs_block = self.context.append_basic_block(func, "land_rhs");
+                    let merge_block = self.context.append_basic_block(func, "land_merge");
+
+                    // If lhs is true, go compute rhs; else jump to merge with false
+                    self.builder
+                        .build_conditional_branch(lhs_bool, rhs_block, merge_block)
+                        .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
+
+                    // RHS path
+                    self.builder.position_at_end(rhs_block);
+                    let rhs_val = self.compile_expr(right)?;
+                    let rhs_int = match rhs_val {
+                        BasicValueEnum::IntValue(iv) => iv,
+                        _ => {
+                            return Err(CodeGenError::TypeError(
+                                "Logical AND requires integer operands".to_string(),
+                            ))
+                        }
+                    };
+                    let rhs_zero = rhs_int.get_type().const_zero();
+                    let rhs_bool = self
+                        .builder
+                        .build_int_compare(
+                            inkwell::IntPredicate::NE,
+                            rhs_int,
+                            rhs_zero,
+                            "land_rhs_nz",
+                        )
+                        .map_err(|e| CodeGenError::Builder(e.to_string()))?;
+                    let rhs_end_block = self.builder.get_insert_block().ok_or_else(|| {
+                        CodeGenError::LLVMError("No current basic block after RHS".to_string())
+                    })?;
+                    self.builder
+                        .build_unconditional_branch(merge_block)
+                        .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
+
+                    // Merge: phi(i1) with false from LHS=false path, RHS bool from rhs path
+                    self.builder.position_at_end(merge_block);
+                    let i1 = self.context.bool_type();
+                    let phi = self
+                        .builder
+                        .build_phi(i1, "land_phi")
+                        .map_err(|e| CodeGenError::LLVMError(e.to_string()))?;
+                    let zero = i1.const_zero();
+                    phi.add_incoming(&[(&rhs_bool, rhs_end_block), (&zero, curr_block)]);
+                    return Ok(phi.as_basic_value());
+                }
+
+                // Default eager evaluation for other binary ops
                 let left_val = self.compile_expr(left)?;
                 let right_val = self.compile_expr(right)?;
                 self.compile_binary_op(left_val, op.clone(), right_val)
diff --git a/ghostscope-compiler/src/script/ast.rs b/ghostscope-compiler/src/script/ast.rs
@@ -212,19 +212,39 @@ pub fn infer_type(expr: &Expr) -> Result<VarType, String> {
                     return Ok(VarType::Bool);
                 }
 
-                // Logical operations expect boolean operands and return boolean
+                // Logical operations return boolean; allow Int literals as truthy (non-zero)
                 if matches!(*op, BinaryOp::LogicalAnd | BinaryOp::LogicalOr) {
-                    if left_type != VarType::Bool || right_type != VarType::Bool {
-                        return Err("Logical operations require boolean operands".to_string());
+                    match (left_type, right_type) {
+                        (VarType::Bool, VarType::Bool)
+                        | (VarType::Bool, VarType::Int)
+                        | (VarType::Int, VarType::Bool)
+                        | (VarType::Int, VarType::Int) => return Ok(VarType::Bool),
+                        _ => {
+                            return Err("Logical operations require boolean or integer operands"
+                                .to_string())
+                        }
                     }
-                    return Ok(VarType::Bool);
                 }
 
                 Ok(left_type)
             } else {
                 // If there are variable references, assume type compatibility to let parsing pass
                 // Actual type checking will be done in code generation phase
-                Ok(VarType::Int)
+                if matches!(*op, BinaryOp::LogicalAnd | BinaryOp::LogicalOr)
+                    || matches!(
+                        *op,
+                        BinaryOp::Equal
+                            | BinaryOp::NotEqual
+                            | BinaryOp::LessThan
+                            | BinaryOp::LessEqual
+                            | BinaryOp::GreaterThan
+                            | BinaryOp::GreaterEqual
+                    )
+                {
+                    Ok(VarType::Bool)
+                } else {
+                    Ok(VarType::Int)
+                }
             }
         }
     }
diff --git a/ghostscope/tests/complex_types_execution.rs b/ghostscope/tests/complex_types_execution.rs
@@ -490,6 +490,87 @@ trace complex_types_program.c:25 {
     Ok(())
 }
 
+#[tokio::test]
+async fn test_or_short_circuit_avoids_null_deref() -> anyhow::Result<()> {
+    init();
+
+    // Start complex_types_program (Debug)
+    let binary_path =
+        FIXTURES.get_test_binary_with_opt("complex_types_program", OptimizationLevel::Debug)?;
+    let mut prog = Command::new(&binary_path)
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .spawn()?;
+    let pid = prog
+        .id()
+        .ok_or_else(|| anyhow::anyhow!("Failed to get PID"))?;
+    tokio::time::sleep(Duration::from_millis(500)).await;
+
+    // The RHS would deref c.friend_ref, which is NULL for 'a' iterations.
+    // Since LHS is true, RHS must not be evaluated and no null-deref error should appear.
+    let script = r#"
+trace update_complex {
+    print (1 || *c.friend_ref);
+}
+"#;
+
+    let (exit_code, stdout, stderr) = run_ghostscope_with_script_for_pid(script, 3, pid).await?;
+    let _ = prog.kill().await;
+    assert_eq!(exit_code, 0, "stderr={} stdout={}", stderr, stdout);
+
+    assert!(
+        stdout.contains("true"),
+        "Expected true result. STDOUT: {}",
+        stdout
+    );
+    assert!(
+        !stdout.contains("<error: null pointer dereference>"),
+        "Short-circuit should avoid null-deref RHS. STDOUT: {}",
+        stdout
+    );
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_and_short_circuit_avoids_null_deref() -> anyhow::Result<()> {
+    init();
+
+    // Start complex_types_program (Debug)
+    let binary_path =
+        FIXTURES.get_test_binary_with_opt("complex_types_program", OptimizationLevel::Debug)?;
+    let mut prog = Command::new(&binary_path)
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .spawn()?;
+    let pid = prog
+        .id()
+        .ok_or_else(|| anyhow::anyhow!("Failed to get PID"))?;
+    tokio::time::sleep(Duration::from_millis(500)).await;
+
+    // LHS is false, RHS would deref c.friend_ref which can be NULL. Short-circuit must avoid RHS.
+    let script = r#"
+trace update_complex {
+    print (0 && *c.friend_ref);
+}
+"#;
+
+    let (exit_code, stdout, stderr) = run_ghostscope_with_script_for_pid(script, 3, pid).await?;
+    let _ = prog.kill().await;
+    assert_eq!(exit_code, 0, "stderr={} stdout={}", stderr, stdout);
+
+    assert!(
+        stdout.contains("false"),
+        "Expected false result. STDOUT: {}",
+        stdout
+    );
+    assert!(
+        !stdout.contains("<error: null pointer dereference>"),
+        "Short-circuit should avoid null-deref RHS. STDOUT: {}",
+        stdout
+    );
+    Ok(())
+}
+
 #[tokio::test]
 async fn test_address_of_and_comparisons_local() -> anyhow::Result<()> {
     init();
diff --git a/ghostscope/tests/script_execution.rs b/ghostscope/tests/script_execution.rs
