loop bound canon lte

hercules_opt/src/loop_bound_canon.rs

@@ -73,6 +73,9 @@ pub fn canonicalize_single_loop_bounds(
         .into_iter()
         .partition(|f| loop_bound_iv_phis.contains(&f.phi()));
 
+    println!("{:?}", loop_bound_ivs);
+
+
     // Assume there is only one loop bound iv.
     if loop_bound_ivs.len() != 1 {
         return false;
@@ -93,9 +96,6 @@ pub fn canonicalize_single_loop_bounds(
         return false;
     };
 
-    let Some(final_value) = final_value else {
-        return false;
-    };
 
     let Some(loop_pred) = editor
         .get_uses(l.header)
@@ -109,8 +109,23 @@ pub fn canonicalize_single_loop_bounds(
 
     // (init_id, bound_id, binop node, if node).
 
+    // FIXME: This is not always correct, depends on lots of things about the loop IV. 
+    let loop_bound_dc = match *editor.node(condition_node) {
+        Node::Binary { left, right, op } => match op {
+            BinaryOperator::LT => right,
+            BinaryOperator::LTE => right,
+            BinaryOperator::GT => {return false}
+            BinaryOperator::GTE => {return false}
+            BinaryOperator::EQ => {return false}
+            BinaryOperator::NE => {return false}
+            _ => {return false}
+        },
+        _ => {return false}
+    };
+
+    
     // FIXME: This is quite fragile.
-    let guard_info: Option<(NodeID, NodeID, NodeID, NodeID)> = (|| {
+    let mut guard_info: Option<(NodeID, NodeID, NodeID, NodeID)> = (|| {
         let Node::ControlProjection {
             control,
             selection: _,
@@ -119,7 +134,7 @@ pub fn canonicalize_single_loop_bounds(
             return None;
         };
 
-        let Node::If { control, cond } = editor.node(control) else {
+        let Node::If { cond, ..} = editor.node(control) else {
             return None;
         };
 
@@ -129,7 +144,7 @@ pub fn canonicalize_single_loop_bounds(
 
         let Node::Binary {
             left: _,
-            right: _,
+            right: r,
             op: loop_op,
         } = editor.node(condition_node)
         else {
@@ -144,7 +159,7 @@ pub fn canonicalize_single_loop_bounds(
             return None;
         }
 
-        if right != final_value {
+        if right != r {
             return None;
         }
 
@@ -169,7 +184,7 @@ pub fn canonicalize_single_loop_bounds(
     // We are assuming this is a simple loop bound (i.e only one induction variable involved), so that .
     let Node::DynamicConstant {
         id: loop_bound_dc_id,
-    } = *editor.node(final_value)
+    } = *editor.node(loop_bound_dc)
     else {
         return false;
     };
@@ -177,9 +192,9 @@ pub fn canonicalize_single_loop_bounds(
     // We need to do 4 (5) things, which are mostly separate.
 
     // 0) Make the update into addition.
-    // 1) Make the update a positive value.
-    // 2) Transform the condition into a `<`
-    // 3) Adjust update to be 1 (and bounds).
+    // 1) Adjust update to be 1 (and bounds).
+    // 2) Make the update a positive value. / Transform the condition into a `<`
+    // - Are these separate?
     // 4) Change init to start from 0.
 
     // 5) Find some way to get fork-guard-elim to work with the new fork.
@@ -198,7 +213,13 @@ pub fn canonicalize_single_loop_bounds(
                     return false;
                 }
             }
-            BinaryOperator::LTE => todo!(),
+            BinaryOperator::LTE => {
+                if left == *update_expression && editor.node(right).is_dynamic_constant() {
+                    right
+                } else {
+                    return false;
+                }
+            }
             BinaryOperator::GT => todo!(),
             BinaryOperator::GTE => todo!(),
             BinaryOperator::EQ => todo!(),
@@ -211,8 +232,10 @@ pub fn canonicalize_single_loop_bounds(
         _ => return false,
     };
 
+    let condition_node_data = editor.node(condition_node).clone();
+
     let Node::DynamicConstant {
-        id: bound_node_dc_id,
+        id: mut bound_node_dc_id,
     } = *editor.node(dc_bound_node)
     else {
         return false;
@@ -220,7 +243,59 @@ pub fn canonicalize_single_loop_bounds(
 
     // If increment is negative (how in the world do we know that...)
     // Increment can be DefinetlyPostiive, Unknown, DefinetlyNegative.
+    let misc_guard_thing: Option<Node> =  if let Some((init_id, bound_id, binop_node, if_node)) = guard_info {
+        Some(editor.node(binop_node).clone())
+    } else {
+        None
+    };
+    println!("condition node: {:?}", condition_node);
+    let users = editor.get_users(condition_node).collect_vec(); 
+    println!("{:?}", users);
+
+    let mut condition_node = condition_node;
+
+    let result = editor.edit(|mut edit| {
+        // 2) Transform the condition into a < (from <=)
+        if let Node::Binary { left, right, op } = condition_node_data {
+            if BinaryOperator::LTE == op && left == *update_expression {
+                // Change the condition into <
+                let new_bop = edit.add_node(Node::Binary { left, right, op: BinaryOperator::LT });
+                
+                // Change the bound dc to be bound_dc + 1
+                let one = DynamicConstant::Constant(1);
+                let one = edit.add_dynamic_constant(one);
+
+                let tmp = DynamicConstant::add(bound_node_dc_id, one);
+                let new_condition_dc = edit.add_dynamic_constant(tmp);
+
+                let new_dc_bound_node = edit.add_node(Node::DynamicConstant { id: new_condition_dc });
+
+                // // 5) Change loop guard:
+                guard_info = if let Some((init_id, bound_id, binop_node, if_node)) = guard_info {
+                    // Change binop node
+                    let Some(Node::Binary { left, right, op }) =  misc_guard_thing else {unreachable!()};
+                    let blah = edit.add_node(Node::DynamicConstant { id:  new_condition_dc});
+
+                    // FIXME: Don't assume that right is the loop bound in the guard. 
+                    let new_binop_node = edit.add_node(Node::Binary { left, right: blah, op: BinaryOperator::LT });
+
+                    edit = edit.replace_all_uses_where(binop_node, new_binop_node, |usee| *usee == if_node)?;
+                    Some((init_id, bound_id, new_binop_node, if_node))
+                } else {guard_info};
+
+                edit = edit.replace_all_uses_where(dc_bound_node, new_dc_bound_node, |usee| *usee == new_bop)?;
+                edit = edit.replace_all_uses(condition_node, new_bop)?;
+
+                // Change loop condition
+                dc_bound_node = new_dc_bound_node;
+                bound_node_dc_id = new_condition_dc;
+                condition_node = new_bop;
+            }
+        };
+        Ok(edit)
+    });
 
+    
     let update_expr_users: Vec<_> = editor
         .get_users(*update_expression)
         .filter(|node| *node != iv.phi() && *node != condition_node)
@@ -241,34 +316,23 @@ pub fn canonicalize_single_loop_bounds(
         let new_init = edit.add_node(new_init);
         edit = edit.replace_all_uses_where(*initializer, new_init, |usee| *usee == iv.phi())?;
 
-        let new_condition_id = DynamicConstant::sub(bound_node_dc_id, init_dc_id);
-        let new_condition = Node::DynamicConstant {
-            id: edit.add_dynamic_constant(new_condition_id),
+        let new_condition_dc = DynamicConstant::sub(bound_node_dc_id, init_dc_id);
+        let new_condition_dc_id = Node::DynamicConstant {
+            id: edit.add_dynamic_constant(new_condition_dc),
         };
-        let new_condition = edit.add_node(new_condition);
+        let new_condition_dc = edit.add_node(new_condition_dc_id);
         edit = edit
-            .replace_all_uses_where(dc_bound_node, new_condition, |usee| *usee == condition_node)?;
+            .replace_all_uses_where(dc_bound_node, new_condition_dc, |usee| *usee == condition_node)?;
 
-        // Change loop guard:
+        // 5) Change loop guard:
         if let Some((init_id, bound_id, binop_node, if_node)) = guard_info {
             edit = edit.replace_all_uses_where(init_id, new_init, |usee| *usee == binop_node)?;
             edit =
-                edit.replace_all_uses_where(bound_id, new_condition, |usee| *usee == binop_node)?;
+                edit.replace_all_uses_where(bound_id, new_condition_dc, |usee| *usee == binop_node)?;
         }
+        
 
-        // for user in update_expr_users {
-        //     let new_user = Node::Binary {
-        //         left: user,
-        //         right: *initializer,
-        //         op: BinaryOperator::Add,
-        //     };
-        //     let new_user = edit.add_node(new_user);
-        //     edit = edit.replace_all_uses(user, new_user)?;
-        // }
-
-        // for
-
-        // Add the offset back to users of the IV update expression
+        // 4) Add the offset back to users of the IV update expression
         let new_user = Node::Binary {
             left: *update_expression,
             right: *initializer,