rewrite group opt

rust2/src/codegen.rs

@@ -19,6 +19,7 @@
 
 use crate::opts::{Ir, Stmt as IrStmt, StmtKind};
 use crate::parse::Span;
+use crate::BumpVec;
 use bumpalo::Bump;
 
 #[derive(Debug, Clone, Copy)]
@@ -39,8 +40,8 @@ const _: [(); 8] = [(); std::mem::size_of::<Stmt>()];
 
 #[derive(Debug, Clone)]
 pub struct Code<'c> {
-    stmts: Vec<Stmt, &'c Bump>,
-    debug: Vec<Span, &'c Bump>,
+    stmts: BumpVec<'c, Stmt>,
+    debug: BumpVec<'c, Span>,
 }
 
 impl Code<'_> {

rust2/src/lib.rs

@@ -13,6 +13,8 @@ pub mod codegen_interpreter;
 pub mod opts;
 pub mod parse;
 
+type BumpVec<'a, T> = Vec<T, &'a Bump>;
+
 pub enum UseProfile {
     Yes,
     No,

rust2/src/opts.rs

@@ -1,9 +1,10 @@
 use crate::parse::{Instr, Span};
+use crate::BumpVec;
 use bumpalo::Bump;
 
 #[derive(Debug, Clone)]
 pub struct Ir<'ir> {
-    pub stmts: Vec<Stmt<'ir>, &'ir Bump>,
+    pub stmts: BumpVec<'ir, Stmt<'ir>>,
 }
 
 #[derive(Debug, Clone)]
@@ -34,64 +35,81 @@ pub enum StmtKind<'ir> {
 }
 
 pub fn optimize<'ir>(alloc: &'ir Bump, instrs: &[(Instr<'_>, Span)]) -> Ir<'ir> {
-    let mut ir = ast_to_ir(alloc, instrs);
+    let ir = ast_to_ir(alloc, instrs);
+    let mut ir = pass_group(alloc, ir);
     pass_find_set_null(&mut ir);
     ir
 }
 
 fn ast_to_ir<'ir>(alloc: &'ir Bump, ast: &[(Instr<'_>, Span)]) -> Ir<'ir> {
     let mut stmts = Vec::new_in(alloc);
-    let mut spans = Vec::new_in(alloc);
 
-    let mut instr_iter = ast.iter();
-
-    let Some(first) = instr_iter.next() else {
-        return Ir { stmts };
-    };
-
-    let mut last = &first.0;
-    let mut start_span = first.1;
-    let mut count = 1;
-    let mut end_span = start_span;
-
-    for (next, next_span) in instr_iter {
-        match last {
-            Instr::Add | Instr::Sub | Instr::Right | Instr::Left if last == next => {
-                count += 1;
-                end_span = *next_span;
-                continue;
+    let stmts_iter = ast.iter().map(|(instr, span)| {
+        let kind = match instr {
+            Instr::Add => StmtKind::Add(1),
+            Instr::Sub => StmtKind::Sub(1),
+            Instr::Right => StmtKind::Right(1),
+            Instr::Left => StmtKind::Left(1),
+            Instr::Out => StmtKind::Out,
+            Instr::In => StmtKind::In,
+            Instr::Loop(body) => {
+                let ir_body = ast_to_ir(alloc, &body);
+                StmtKind::Loop(ir_body)
             }
-            _ => {
-                end_span = *next_span;
-                let new_last = match last {
-                    Instr::Add => Stmt::lol(StmtKind::Add(count)),
-                    Instr::Sub => Stmt::lol(StmtKind::Sub(count)),
-                    Instr::Right => Stmt::lol(StmtKind::Right(count.into())),
-                    Instr::Left => Stmt::lol(StmtKind::Left(count.into())),
-                    Instr::Out => Stmt::lol(StmtKind::Out),
-                    Instr::In => Stmt::lol(StmtKind::In),
-                    Instr::Loop(body) => Stmt::lol(StmtKind::Loop(ast_to_ir(alloc, body))),
-                };
-                stmts.push(new_last);
-                spans.push(start_span.until(end_span));
-                last = next;
-                start_span = *next_span;
-                count = 1;
-            }
-        }
-    }
+        };
+        Stmt { kind, span: *span }
+    });
 
-    let new_last = match last {
-        Instr::Add => Stmt::lol(StmtKind::Add(count)),
-        Instr::Sub => Stmt::lol(StmtKind::Sub(count)),
-        Instr::Right => Stmt::lol(StmtKind::Right(count.into())),
-        Instr::Left => Stmt::lol(StmtKind::Left(count.into())),
-        Instr::Out => Stmt::lol(StmtKind::Out),
-        Instr::In => Stmt::lol(StmtKind::In),
-        Instr::Loop(body) => Stmt::lol(StmtKind::Loop(ast_to_ir(alloc, &body))),
-    };
-    stmts.push(new_last);
-    spans.push(start_span.until(end_span));
+    stmts.extend(stmts_iter);
+
+    Ir { stmts }
+}
+
+fn pass_group<'ir>(alloc: &'ir Bump, ir: Ir<'ir>) -> Ir<'ir> {
+    let new_stmts = Vec::new_in(alloc);
+    let stmts = ir
+        .stmts
+        .into_iter()
+        .fold(new_stmts, |mut stmts: BumpVec<'ir, Stmt<'ir>>, next| {
+            let Some(old) = stmts.last_mut() else {
+                stmts.push(next);
+                return stmts;
+            };
+
+            match (&mut old.kind, next.kind) {
+                (StmtKind::Add(a), StmtKind::Add(b)) => {
+                    old.span = old.span.merge(next.span);
+                    *a += b;
+                }
+                (StmtKind::Sub(a), StmtKind::Sub(b)) => {
+                    old.span = old.span.merge(next.span);
+                    *a += b;
+                }
+                (StmtKind::Right(a), StmtKind::Right(b)) => {
+                    old.span = old.span.merge(next.span);
+                    *a += b;
+                }
+                (StmtKind::Left(a), StmtKind::Left(b)) => {
+                    old.span = old.span.merge(next.span);
+                    *a += b;
+                }
+                (_, StmtKind::Loop(body)) => {
+                    let new_body = pass_group(alloc, body);
+                    stmts.push(Stmt {
+                        span: next.span,
+                        kind: StmtKind::Loop(new_body),
+                    });
+                }
+                (_, kind) => {
+                    stmts.push(Stmt {
+                        span: next.span,
+                        kind,
+                    });
+                }
+            }
+
+            stmts
+        });
 
     Ir { stmts }
 }

rust2/src/parse.rs

@@ -1,4 +1,5 @@
 use bumpalo::Bump;
+use std::cmp;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
 pub struct Span {
@@ -21,6 +22,7 @@ impl Span {
         }
     }
 
+    #[must_use]
     pub fn until(&self, other: Self) -> Self {
         Self {
             start: self.start,
@@ -28,6 +30,14 @@ impl Span {
         }
     }
 
+    #[must_use]
+    pub fn merge(&self, other: Self) -> Self {
+        Self::start_end(
+            cmp::min(self.start(), other.start()),
+            cmp::max(self.end(), other.end()),
+        )
+    }
+
     pub fn start(&self) -> usize {
         self.start.try_into().unwrap()
     }
@@ -35,6 +45,10 @@ impl Span {
     pub fn len(&self) -> usize {
         self.len.try_into().unwrap()
     }
+
+    pub fn end(&self) -> usize {
+        self.start() + self.len() - 1
+    }
 }
 
 pub type Ast<'ast> = Vec<(Instr<'ast>, Span), &'ast Bump>;