add opt that doesn't make anything faster but might allow better opts later

rust2/src/opts.rs

@@ -14,11 +14,12 @@ pub struct Stmt<'ir> {
 }
 
 impl<'ir> Stmt<'ir> {
-    fn lol(kind: StmtKind<'ir>) -> Stmt<'ir> {
-        Self {
-            kind,
-            span: Span::default(),
-        }
+    fn new(kind: StmtKind<'ir>, span: Span) -> Stmt<'ir> {
+        Self { kind, span }
+    }
+
+    fn kind(&self) -> &StmtKind<'ir> {
+        &self.kind
     }
 }
 
@@ -32,12 +33,14 @@ pub enum StmtKind<'ir> {
     Out,
     In,
     SetNull,
+    SetN(u8),
 }
 
 pub fn optimize<'ir>(alloc: &'ir Bump, instrs: &[(Instr<'_>, Span)]) -> Ir<'ir> {
     let ir = ast_to_ir(alloc, instrs);
     let mut ir = pass_group(alloc, ir);
     pass_find_set_null(&mut ir);
+    pass_set_n(&mut ir);
     ir
 }
 
@@ -57,7 +60,7 @@ fn ast_to_ir<'ir>(alloc: &'ir Bump, ast: &[(Instr<'_>, Span)]) -> Ir<'ir> {
                 StmtKind::Loop(ir_body)
             }
         };
-        Stmt { kind, span: *span }
+        Stmt::new(kind, *span)
     });
 
     stmts.extend(stmts_iter);
@@ -65,6 +68,7 @@ fn ast_to_ir<'ir>(alloc: &'ir Bump, ast: &[(Instr<'_>, Span)]) -> Ir<'ir> {
     Ir { stmts }
 }
 
+/// pass that replaces things like `Sub(1) Sub(1)` with `Sub(2)`
 fn pass_group<'ir>(alloc: &'ir Bump, ir: Ir<'ir>) -> Ir<'ir> {
     let new_stmts = Vec::new_in(alloc);
     let stmts = ir
@@ -74,10 +78,10 @@ fn pass_group<'ir>(alloc: &'ir Bump, ir: Ir<'ir>) -> Ir<'ir> {
             let Some(old) = stmts.last_mut() else {
                 if let StmtKind::Loop(body) = next.kind {
                     let new_body = pass_group(alloc, body);
-                    stmts.push(Stmt {
-                        span: next.span,
-                        kind: StmtKind::Loop(new_body),
-                    });
+                    stmts.push(Stmt::new(
+                         StmtKind::Loop(new_body),
+                        next.span,
+                    ));
                 } else {
                     stmts.push(next);
                 }
@@ -109,10 +113,7 @@ fn pass_group<'ir>(alloc: &'ir Bump, ir: Ir<'ir>) -> Ir<'ir> {
                     });
                 }
                 (_, kind) => {
-                    stmts.push(Stmt {
-                        span: next.span,
-                        kind,
-                    });
+                    stmts.push(Stmt::new(kind, next.span));
                 }
             }
 
@@ -122,6 +123,7 @@ fn pass_group<'ir>(alloc: &'ir Bump, ir: Ir<'ir>) -> Ir<'ir> {
     Ir { stmts }
 }
 
+/// pass that replaces `Loop([Sub(_)])` to `SetNull`
 fn pass_find_set_null(ir: &mut Ir<'_>) {
     for stmt in &mut ir.stmts {
         if let Stmt {
@@ -131,13 +133,37 @@ fn pass_find_set_null(ir: &mut Ir<'_>) {
         {
             if let [Stmt {
                 kind: StmtKind::Sub(_),
-                ..
+                span,
             }] = body.stmts.as_slice()
             {
-                *stmt = Stmt::lol(StmtKind::SetNull);
+                *stmt = Stmt::new(StmtKind::SetNull, *span);
             } else {
                 pass_find_set_null(body);
             }
         }
     }
 }
+
+/// pass that replaces `SetNull Add(5)` with `SetN(5)`
+fn pass_set_n(ir: &mut Ir<'_>) {
+    let stmts = &mut ir.stmts;
+    let mut i = 0;
+    while i < stmts.len() - 1 {
+        let a = &stmts[i];
+        if let StmtKind::SetNull = a.kind() {
+            let b = &stmts[i + 1];
+            let new = match b.kind() {
+                StmtKind::Add(n) => StmtKind::SetN(*n),
+                StmtKind::Sub(n) => StmtKind::SetN(0u8.wrapping_sub(*n)),
+                _ => {
+                    i += 1;
+                    continue;
+                }
+            };
+            let span = a.span.merge(b.span);
+            stmts.remove(i + 1);
+            stmts[i] = Stmt::new(new, span);
+        }
+        i += 1;
+    }
+}