Allow obtaining BasicBlock from get_operand

Author: TheDan64

src/values/basic_value_use.rs

@@ -144,7 +144,7 @@ impl BasicValueUse {
     /// let free_instruction = builder.build_free(arg1);
     /// let return_instruction = builder.build_return(None);
     ///
-    /// let free_operand0 = free_instruction.get_operand(0).unwrap();
+    /// let free_operand0 = free_instruction.get_operand(0).unwrap().left().unwrap();
     /// let free_operand0_instruction = free_operand0.as_instruction_value().unwrap();
     /// let bitcast_use_value = free_operand0_instruction
     ///     .get_first_use()

src/values/instruction_value.rs

@@ -1,4 +1,5 @@
-use llvm_sys::core::{LLVMGetInstructionOpcode, LLVMIsTailCall, LLVMGetPreviousInstruction, LLVMGetNextInstruction, LLVMGetInstructionParent, LLVMInstructionEraseFromParent, LLVMInstructionClone, LLVMSetVolatile, LLVMGetVolatile, LLVMGetNumOperands, LLVMGetOperand, LLVMGetOperandUse, LLVMSetOperand};
+use either::{Either, Either::{Left, Right}};
+use llvm_sys::core::{LLVMGetInstructionOpcode, LLVMIsTailCall, LLVMGetPreviousInstruction, LLVMGetNextInstruction, LLVMGetInstructionParent, LLVMInstructionEraseFromParent, LLVMInstructionClone, LLVMSetVolatile, LLVMGetVolatile, LLVMGetNumOperands, LLVMGetOperand, LLVMGetOperandUse, LLVMSetOperand, LLVMValueAsBasicBlock, LLVMIsABasicBlock};
 #[llvm_versions(3.9 => latest)]
 use llvm_sys::core::LLVMInstructionRemoveFromParent;
 use llvm_sys::LLVMOpcode;
@@ -464,7 +465,7 @@ impl InstructionValue {
     /// 3) Function call has two: i8 pointer %1 argument, and the free function itself
     /// 4) Void return has zero: void is not a value and does not count as an operand
     /// even though the return instruction can take values.
-    pub fn get_operand(&self, index: u32) -> Option<BasicValueEnum> {
+    pub fn get_operand(&self, index: u32) -> Option<Either<BasicValueEnum, BasicBlock>> {
         let num_operands = self.get_num_operands();
 
         if index >= num_operands {
@@ -479,7 +480,19 @@ impl InstructionValue {
             return None;
         }
 
-        Some(BasicValueEnum::new(operand))
+        let is_basic_block = unsafe {
+            !LLVMIsABasicBlock(operand).is_null()
+        };
+
+        if is_basic_block {
+            let operand = unsafe {
+                LLVMValueAsBasicBlock(operand)
+            };
+
+            Some(Right(BasicBlock::new(operand).expect("BasicBlock should be valid")))
+        } else {
+            Some(Left(BasicValueEnum::new(operand)))
+        }
     }
 
     /// Sets the operand an `InstructionValue` has at a given index if possible.
@@ -511,7 +524,7 @@ impl InstructionValue {
     /// // This will produce invalid IR:
     /// free_instruction.set_operand(0, f32_val);
     ///
-    /// assert_eq!(free_instruction.get_operand(0).unwrap(), f32_val);
+    /// assert_eq!(free_instruction.get_operand(0).unwrap().left().unwrap(), f32_val);
     /// ```
     pub fn set_operand<BV: BasicValue>(&self, index: u32, val: BV) -> bool {
         let num_operands = self.get_num_operands();

tests/all/test_instruction_values.rs

@@ -38,20 +38,20 @@ fn test_operands() {
     let store_operand0 = store_instruction.get_operand(0).unwrap();
     let store_operand1 = store_instruction.get_operand(1).unwrap();
 
-    assert_eq!(store_operand0, f32_val); // f32 const
-    assert_eq!(store_operand1, arg1); // f32* arg1
+    assert_eq!(store_operand0.left().unwrap(), f32_val); // f32 const
+    assert_eq!(store_operand1.left().unwrap(), arg1); // f32* arg1
     assert!(store_instruction.get_operand(2).is_none());
     assert!(store_instruction.get_operand(3).is_none());
     assert!(store_instruction.get_operand(4).is_none());
 
-    let free_operand0 = free_instruction.get_operand(0).unwrap();
-    let free_operand1 = free_instruction.get_operand(1).unwrap();
+    let free_operand0 = free_instruction.get_operand(0).unwrap().left().unwrap();
+    let free_operand1 = free_instruction.get_operand(1).unwrap().left().unwrap();
     let free_operand0_instruction = free_operand0.as_instruction_value().unwrap();
 
     assert!(free_operand0.is_pointer_value()); // (implictly casted) i8* arg1
     assert!(free_operand1.is_pointer_value()); // Free function ptr
     assert_eq!(free_operand0_instruction.get_opcode(), BitCast);
-    assert_eq!(free_operand0_instruction.get_operand(0).unwrap(), arg1);
+    assert_eq!(free_operand0_instruction.get_operand(0).unwrap().left().unwrap(), arg1);
     assert!(free_operand0_instruction.get_operand(1).is_none());
     assert!(free_operand0_instruction.get_operand(2).is_none());
     assert!(free_instruction.get_operand(2).is_none());
@@ -84,7 +84,7 @@ fn test_operands() {
         .get_first_use()
         .unwrap()
         .get_used_value();
-    let free_call_param = free_instruction.get_operand(0).unwrap();
+    let free_call_param = free_instruction.get_operand(0).unwrap().left().unwrap();
 
     assert_eq!(bitcast_use_value, free_call_param);
 
@@ -132,6 +132,30 @@ fn test_operands() {
     assert!(module.verify().is_ok());
 }
 
+#[test]
+fn test_basic_block_operand() {
+    let context = Context::create();
+    let module = context.create_module("ivs");
+    let builder = context.create_builder();
+    let void_type = context.void_type();
+    let fn_type = void_type.fn_type(&[], false);
+    let function = module.add_function("bb_op", fn_type, None);
+    let basic_block = context.append_basic_block(&function, "entry");
+    let basic_block2 = context.append_basic_block(&function, "exit");
+
+    builder.position_at_end(&basic_block);
+
+    let branch_instruction = builder.build_unconditional_branch(&basic_block2);
+    let bb_operand = branch_instruction.get_operand(0).unwrap().right().unwrap();
+
+    assert_eq!(bb_operand, basic_block2);
+
+    builder.position_at_end(&basic_block2);
+    builder.build_return(None);
+
+    assert!(module.verify().is_ok());
+}
+
 #[test]
 fn test_get_next_use() {
     let context = Context::create();