Refactor

rikhuijzer · Sep 27, 2024 · f67727c · f67727c
1 parent 7e65cb0
commit f67727c
Show file tree

Hide file tree

Showing 6 changed files with 60 additions and 34 deletions.
diff --git a/src/dialect/arith/op.rs b/src/dialect/arith/op.rs
@@ -6,6 +6,7 @@ use crate::ir::Op;
 use crate::ir::OpOperand;
 use crate::ir::OpResult;
 use crate::ir::OperationName;
+use crate::ir::OperationOperands;
 use crate::ir::Type;
 use crate::ir::Value;
 use crate::parser::Parse;
@@ -59,19 +60,22 @@ impl Parse for ConstantOp {
         let operation_name = parser.expect(TokenKind::BareIdentifier)?;
         assert!(operation_name.lexeme == "arith.constant");
         operation.set_name(ConstantOp::operation_name());
-        let operand = match parser.peek().kind {
+        let operand: OpOperand = match parser.peek().kind {
             TokenKind::Integer => {
                 let integer = parser.advance();
                 let integer = integer.lexeme.clone();
                 let typ = parser.typ()?;
-                let argument = BlockArgument::new(integer, typ);
-                Value::BlockArgument(argument)
+                // Determine value.
+                todo!();
+                // let operand = OpOperand::new(
+                // Value::BlockArgument(argument)
             }
             _ => {
                 return Err(anyhow::anyhow!("Expected integer constant"));
             }
         };
-        operation.set_operands(Arc::new(vec![Arc::new(operand)]));
+        let operand = Arc::new(RwLock::new(operand));
+        operation.set_operands(Arc::new(RwLock::new(vec![operand])));
         operation.set_parent(parent);
         let operation = Arc::new(RwLock::new(operation));
         Ok(Arc::new(RwLock::new(ConstantOp { operation })))
@@ -86,6 +90,7 @@ impl AddiOp {
     /// Canonicalize `addi(addi(x, c0), c1) -> addi(x, c0 + c1)`.
     fn addi_add_constant(&mut self) -> CanonicalizeResult {
         let operands = self.operation().read().unwrap().operands();
+        let operands = operands.read().unwrap();
         assert!(operands.len() == 2);
         let lhs = &operands[0];
         let rhs = &operands[1];
@@ -112,25 +117,27 @@ impl Op for AddiOp {
 }
 
 impl<T: Parse> Parser<T> {
-    fn operand(&mut self, parent: Arc<RwLock<Block>>) -> Result<Arc<OpOperand>> {
+    fn operand(&mut self, parent: Arc<RwLock<Block>>) -> Result<Arc<RwLock<OpOperand>>> {
         let identifier = self.expect(TokenKind::PercentIdentifier)?;
         let name = identifier.lexeme.clone();
         let block = parent.read().unwrap();
-        let first_use = block.first_use(name.clone());
-        let typ = Type::new("any".to_string());
-        let value = Value::OpResult(OpResult::new(name, typ));
-        let operand = OpOperand::new(value, name);
-        Ok(Arc::new(operand))
-    }
-    pub fn operands(&mut self, parent: Arc<RwLock<Block>>) -> Result<Vec<OpOperand>> {
+        let assignment = block.assignment(name.clone());
+        todo!();
+        // let typ = Type::new("any".to_string());
+        // let value = Value::OpResult(OpResult::new(name, typ));
+        // let operand = OpOperand::new(value, name);
+        // Ok(Arc::new(operand))
+    }
+    pub fn operands(&mut self, parent: Arc<RwLock<Block>>) -> Result<OperationOperands> {
         let mut arguments = vec![];
         while self.check(TokenKind::PercentIdentifier) {
-            arguments.push(self.operand(parent.clone())?);
+            let operand = self.operand(parent.clone())?;
+            arguments.push(operand);
             if self.check(TokenKind::Comma) {
                 let _comma = self.advance();
             }
         }
-        Ok(arguments)
+        Ok(Arc::new(RwLock::new(arguments)))
     }
     pub fn results(&mut self) -> Result<Vec<Arc<Value>>> {
         let mut results = vec![];
@@ -161,7 +168,7 @@ impl Parse for AddiOp {
         operation.set_name(AddiOp::operation_name());
         assert!(parent.is_some());
         operation.set_parent(parent.clone());
-        operation.set_operands(Arc::new(parser.operands(parent.unwrap())?));
+        operation.set_operands(parser.operands(parent.unwrap())?);
         let _colon = parser.expect(TokenKind::Colon)?;
         let result_type = parser.expect(TokenKind::IntType)?;
         let result_type = Type::new(result_type.lexeme.clone());

diff --git a/src/dialect/func/op.rs b/src/dialect/func/op.rs
@@ -3,6 +3,7 @@ use crate::ir::BlockArgument;
 use crate::ir::Op;
 use crate::ir::OpResult;
 use crate::ir::Operation;
+use crate::ir::OperationArguments;
 use crate::ir::OperationName;
 use crate::ir::Type;
 use crate::ir::Value;
@@ -39,8 +40,10 @@ impl Op for FuncOp {
             .read()
             .unwrap()
             .operands()
+            .read()
+            .unwrap()
             .iter()
-            .map(|o| o.to_string())
+            .map(|o| o.read().unwrap().to_string())
             .collect::<Vec<String>>()
             .join(", ");
         write!(f, "{}", joined)?;
@@ -85,7 +88,7 @@ impl<T: Parse> Parser<T> {
         Ok(identifier.lexeme.clone())
     }
     /// %arg0 : i64,
-    pub fn function_argument(&mut self) -> Result<Arc<Value>> {
+    pub fn function_argument(&mut self) -> Result<Arc<RwLock<BlockArgument>>> {
         let identifier = self.expect(TokenKind::PercentIdentifier)?;
         let name = identifier.lexeme.clone();
         let typ = if self.check(TokenKind::Colon) {
@@ -96,14 +99,14 @@ impl<T: Parse> Parser<T> {
             Type::new("any".to_string())
         };
         let arg = BlockArgument::new(name, typ);
-        let operand: Value = Value::BlockArgument(arg);
+        let operand = Arc::new(RwLock::new(arg));
         if self.check(TokenKind::Comma) {
             self.advance();
         }
-        Ok(Arc::new(operand))
+        Ok(operand)
     }
     /// Parse `(%arg0 : i64, %arg1 : i64)`.
-    pub fn function_arguments(&mut self) -> Result<Vec<Arc<Value>>> {
+    pub fn function_arguments(&mut self) -> Result<OperationArguments> {
         let _lparen = self.expect(TokenKind::LParen)?;
         let mut operands = vec![];
         while self.check(TokenKind::PercentIdentifier) {
@@ -115,7 +118,7 @@ impl<T: Parse> Parser<T> {
         } else {
             return Err(anyhow::anyhow!("Expected ')', got {:?}", self.peek().kind));
         }
-        Ok(operands)
+        Ok(Arc::new(RwLock::new(operands)))
     }
     pub fn result_types(&mut self) -> Result<Vec<Type>> {
         let mut result_types = vec![];
@@ -150,7 +153,7 @@ impl Parse for FuncOp {
         let _operation_name = parser.advance();
         let identifier = parser.identifier(TokenKind::AtIdentifier).unwrap();
         let mut operation = Operation::default();
-        operation.set_operands(Arc::new(parser.function_arguments()?));
+        operation.set_arguments(parser.function_arguments()?);
         operation.set_result_types(parser.result_types()?);
         operation.set_region(parser.region()?);
         operation.set_parent(parent);
@@ -186,8 +189,9 @@ impl Op for ReturnOp {
         let operation = operation.read().unwrap();
         let operands = operation.operands().clone();
         write!(f, "return")?;
-        for operand in &*operands {
-            write!(f, " {}", operand)?;
+        let operands = operands.read().unwrap();
+        for operand in operands.iter() {
+            write!(f, " {}", operand.read().unwrap())?;
         }
         let result_types = operation.result_types();
         assert!(!result_types.is_empty(), "Expected result types to be set");
@@ -211,7 +215,7 @@ impl Parse for ReturnOp {
         let mut operation = Operation::default();
         assert!(parent.is_some());
         operation.set_parent(parent.clone());
-        operation.set_operands(Arc::new(parser.operands(parent.clone().unwrap())?));
+        operation.set_operands(parser.operands(parent.clone().unwrap())?);
         let _colon = parser.expect(TokenKind::Colon)?;
         let return_type = parser.expect(TokenKind::IntType)?;
         let return_type = Type::new(return_type.lexeme.clone());

diff --git a/src/ir/block.rs b/src/ir/block.rs
@@ -49,7 +49,7 @@ impl Block {
             for operand in operands.iter() {
                 let other_operand = operand.clone();
                 let write = other_operand.write().unwrap();
-                let value = operand.write().unwrap().value();
+                let value = other_operand.write().unwrap().value();
                 let value = value.clone().read().unwrap();
                 match value {
                     Value::BlockArgument(block_argument) => {

diff --git a/src/ir/mod.rs b/src/ir/mod.rs
@@ -12,7 +12,9 @@ pub use crate::ir::block::Block;
 pub use crate::ir::module::ModuleOp;
 pub use crate::ir::op::Op;
 pub use crate::ir::operation::Operation;
+pub use crate::ir::operation::OperationArguments;
 pub use crate::ir::operation::OperationName;
+pub use crate::ir::operation::OperationOperands;
 pub use crate::ir::region::Region;
 pub use crate::ir::value::BlockArgument;
 pub use crate::ir::value::OpOperand;

diff --git a/src/ir/operation.rs b/src/ir/operation.rs
@@ -25,6 +25,9 @@ impl OperationName {
     }
 }
 
+pub type OperationArguments = Arc<RwLock<Vec<Arc<RwLock<BlockArgument>>>>>;
+pub type OperationOperands = Arc<RwLock<Vec<Arc<RwLock<OpOperand>>>>>;
+
 /// Note that MLIR distinguishes between Operation and Op.
 /// Operation generically models all operations.
 /// Op is an interface for more specific operations.
@@ -39,8 +42,8 @@ impl OperationName {
 pub struct Operation {
     name: OperationName,
     /// Used by `FuncOp` to store its arguments.
-    arguments: Arc<RwLock<Vec<Arc<RwLock<BlockArgument>>>>>,
-    operands: Arc<RwLock<Vec<Arc<RwLock<OpOperand>>>>>,
+    arguments: OperationArguments,
+    operands: OperationOperands,
     attributes: Vec<Arc<dyn Attribute>>,
     /// Results can be `Values`, so either `BlockArgument` or `OpResult`.
     results: Vec<Arc<Value>>,
@@ -55,8 +58,8 @@ pub struct Operation {
 impl Operation {
     pub fn new(
         name: OperationName,
-        arguments: Arc<RwLock<Vec<Arc<RwLock<BlockArgument>>>>>,
-        operands: Arc<RwLock<Vec<Arc<RwLock<OpOperand>>>>>,
+        arguments: OperationArguments,
+        operands: OperationOperands,
         attributes: Vec<Arc<dyn Attribute>>,
         results: Vec<Arc<Value>>,
         result_types: Vec<Type>,
@@ -77,13 +80,13 @@ impl Operation {
     pub fn name(&self) -> String {
         self.name.name.clone()
     }
-    pub fn arguments(&self) -> Arc<RwLock<Vec<Arc<RwLock<BlockArgument>>>>> {
+    pub fn arguments(&self) -> OperationArguments {
         self.arguments.clone()
     }
-    pub fn operands(&self) -> Arc<RwLock<Vec<Arc<RwLock<OpOperand>>>>> {
+    pub fn operands(&self) -> OperationOperands {
         self.operands.clone()
     }
-    pub fn operands_mut(&mut self) -> &mut Arc<RwLock<Vec<Arc<RwLock<OpOperand>>>>> {
+    pub fn operands_mut(&mut self) -> &mut OperationOperands {
         &mut self.operands
     }
     pub fn attributes(&self) -> Vec<Arc<dyn Attribute>> {
@@ -105,7 +108,10 @@ impl Operation {
     pub fn set_name(&mut self, name: OperationName) {
         self.name = name;
     }
-    pub fn set_operands(&mut self, operands: Arc<RwLock<Vec<Arc<RwLock<OpOperand>>>>>) {
+    pub fn set_arguments(&mut self, arguments: OperationArguments) {
+        self.arguments = arguments;
+    }
+    pub fn set_operands(&mut self, operands: OperationOperands) {
         self.operands = operands;
     }
     pub fn set_attributes(&mut self, attributes: Vec<Arc<dyn Attribute>>) {

diff --git a/src/ir/value.rs b/src/ir/value.rs
@@ -70,6 +70,7 @@ impl Display for Value {
         }
     }
 }
+
 pub struct OpOperand {
     pub value: Arc<RwLock<Value>>,
     pub operand_name: String,
@@ -89,3 +90,9 @@ impl OpOperand {
         &self.operand_name
     }
 }
+
+impl Display for OpOperand {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self)
+    }
+}