naive wadler impl

lib.rs

use std::collections::VecDeque;

#[derive(Clone, Debug)]
pub enum Doc {
    Nil,
    Concat(Box<Doc>, Box<Doc>),
    Nest(i32, Box<Doc>),
    Text(String),
    Line,
    Union(Box<Doc>, Box<Doc>),
}

pub enum DocToken<'a> {
    Nil,
    Text(&'a str),
    Line(i32)
}

impl Doc {
    pub fn nil() -> Self {
        Doc::Nil
    }

    pub fn concat(x: Doc, y: Doc) -> Doc {
        Doc::Concat(Box::new(x), Box::new(y))
    }

    pub fn nest(i: i32, x: Doc) -> Doc {
        Doc::Nest(i, Box::new(x))
    }

    pub fn text(s: String) -> Doc {
        Doc::Text(s)
    }

    pub fn line() -> Doc {
        Doc::Line
    }

    pub fn group(x: Doc) -> Doc {
        Doc::union(flatten(x.clone()), x)
    }

    fn union(x: Doc, y: Doc) -> Doc {
        Doc::Union(Box::new(x), Box::new(y))
    }
}

fn flatten(d: Doc) -> Doc {
    match d {
        Doc::Nil => Doc::Nil,
        Doc::Concat(x, y) => Doc::Concat(Box::new(flatten(*x)), Box::new(flatten(*y))),
        Doc::Nest(i, x) => Doc::Nest(i, Box::new(flatten(*x))),
        Doc::Text(s) => Doc::Text(s),
        Doc::Line => Doc::Text(" ".to_owned()),
        Doc::Union(x, y) => flatten(*x),
    }
}

pub fn layout(ts: Vec<DocToken>) -> String {
    let size = ts.iter().fold(0, |mut acc, t| {
        match t {
            DocToken::Nil => 0,
            DocToken::Text(s) => s.len(),
            DocToken::Line(i) => *i as usize + 1,
        }
    });
    let mut buffer = String::with_capacity(size);
    for t in ts.into_iter() {
        match t {
            DocToken::Nil => (),
            DocToken::Text(s) => buffer.push_str(s),
            DocToken::Line(i) => buffer.push_str(&format!("\n{:width$}", "", width = i as usize)),
        }
    }
    
    buffer
}

fn best<'a>(w: i32, k: i32, x: &'a Doc) -> Vec<DocToken<'a>> {
    be(w, k, vec![(0, &x)])
}

fn be<'a>(w: i32, mut k: i32, mut z: Vec<(i32, &'a Doc)>) -> Vec<DocToken<'a>> {
    if z.is_empty() {
        return vec![DocToken::Nil];
    }
    let mut tokens = Vec::new();
    while let Some((i, doc)) = z.pop() {
        match doc {
            Doc::Nil => (),
            Doc::Concat(x, y) => {
                z.push((i, y));
                z.push((i, x));
            }
            Doc::Nest(j, x) => {
                z.push((i + j, x));
            }
            Doc::Text(ref s) => {
                let length = s.len() as i32;
                tokens.push(DocToken::Text(s));
                k += length;
            }
            Doc::Line => {
                tokens.push(DocToken::Line(i));
                k = i;
            },
            Doc::Union(x, y) => {
                let mut z_left = z.clone();
                z_left.push((i, x));
                let mut z_right = z;
                z_right.push((i, y));
                tokens.extend(better(w, k, be(w, k, z_left), be(w, k, z_right)));
                return tokens;
            }
        }
    }

    tokens
}

fn better<'a>(w: i32, k: i32, x: Vec<DocToken<'a>>, y: Vec<DocToken<'a>>) -> Vec<DocToken<'a>> {
    if fits(w - k, &x) {
        x
    } else {
        y
    }
}

fn fits(mut w: i32, ts: &Vec<DocToken>) -> bool {
    if w < 0 {
        return false;
    }
    for t in ts {
        match t {
            DocToken::Nil => (),
            DocToken::Text(s) => if s.len() as i32 > w {
                return false;
            } else {
                w -= s.len() as i32;
            }
            DocToken::Line(_) => return true,
        }
    }
    return true;
}

pub fn pretty(w: i32, x: Doc) -> String {
    layout(best(w, 0, &x))
}

//// Utility

// x <> text " " <> y
pub fn space_join(x: Doc, y: Doc) -> Doc {
    Doc::concat(x, Doc::concat(Doc::Text(" ".to_owned()), y))
}

// x <> line <> y
pub fn newline_join(x: Doc, y: Doc) -> Doc {
    Doc::concat(x, Doc::concat(Doc::Line, y))
}

pub fn folddoc(f: impl Fn(Doc, Doc) -> Doc, mut xs: Vec<Doc>) -> Doc {
    if xs.is_empty() {
        return Doc::Nil;
    }
    if xs.len() == 1 {
        return xs.pop().unwrap();
    }
    xs.into_iter()
        .rev()
        .fold(Doc::Nil, |acc, item| f(item, acc))
}

pub fn spread(xs: Vec<Doc>) -> Doc {
    folddoc(space_join, xs)
}

pub fn stack(xs: Vec<Doc>) -> Doc {
    folddoc(newline_join, xs)
}

pub fn bracket(l: String, x: Doc, r: String) -> Doc {
    Doc::group(Doc::concat(
        Doc::text(l),
        Doc::concat(
            Doc::nest(2, Doc::concat(Doc::Line, x)),
            Doc::concat(Doc::Line, Doc::text(r)),
        ),
    ))
}

pub fn break_join(x: Doc, y: Doc) -> Doc {
    Doc::concat(
        x,
        Doc::concat(Doc::union(Doc::text(" ".to_owned()), Doc::Line), y),
    )
}

pub fn fillwords(s: String) -> Doc {
    folddoc(
        break_join,
        s.split_whitespace()
            .map(|x| Doc::text(x.to_owned()))
            .collect(),
    )
}

pub fn fill(zs: Vec<Doc>) -> Doc {
    let mut q = VecDeque::with_capacity(zs.len());
    q.extend(zs.into_iter());
    fill_inner(q)
}

    
fn fill_inner(mut zs: VecDeque<Doc>) -> Doc {
    if zs.is_empty() {
        return Doc::Nil;
    }
    if zs.len() == 1 {
        return zs.pop_front().unwrap();
    }

    let x = zs.pop_front().unwrap();
    let mut yzs = zs.clone();
    let y = zs.pop_front().unwrap();

    Doc::union(
        flatten(space_join(
            x.clone(),
            fill_inner({
                zs.push_front(flatten(y));
                zs
            }),
        )),
        newline_join(x, fill_inner(yzs)),
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::VecDeque;

    // Tree example
    #[derive(Debug, Clone)]
    enum Tree {
        Node(&'static str, VecDeque<Tree>),
    }

    macro_rules! list{
        ( $($e:expr),* ) => {
            {
                let v = vec![ $($e),* ];
                let mut q = VecDeque::with_capacity(v.len());
                q.extend(v.into_iter());
                q
            }
        }
    }

    fn show_tree(Tree::Node(s, ts): Tree) -> Doc {
        let length = s.len() as i32;
        Doc::group(Doc::concat(
            Doc::text(s.to_owned()),
            Doc::nest(length, show_bracket(ts)),
        ))
    }

    fn show_bracket(ts: VecDeque<Tree>) -> Doc {
        if ts.len() == 0 {
            return Doc::Nil;
        }
        Doc::concat(
            Doc::text("[".to_owned()),
            Doc::concat(Doc::nest(1, show_trees(ts)), Doc::text("]".to_owned())),
        )
    }

    fn show_trees(mut ts: VecDeque<Tree>) -> Doc {
        assert!(ts.len() > 0);
        let t = ts.pop_front().unwrap();
        if ts.len() == 0 {
            return show_tree(t);
        }
        Doc::concat(
            show_tree(t),
            Doc::concat(
                Doc::text(",".to_owned()),
                Doc::concat(Doc::Line, show_trees(ts)),
            ),
        )
    }

    #[test]
    fn print_tree() {
        let tree = Tree::Node(
            "aaa",
            list![
                Tree::Node(
                    "bbbbb",
                    list![Tree::Node("ccc", list![]), Tree::Node("dd", list![])]
                ),
                Tree::Node("eee", list![]),
                Tree::Node(
                    "ffff",
                    list![
                        Tree::Node("gg", list![]),
                        Tree::Node("hhh", list![]),
                        Tree::Node("ii", list![])
                    ]
                )
            ],
        );
        let doc = show_tree(tree.clone());
        // println!("{:#?}", doc);
        println!("{}", pretty(30, doc.clone()));
        let expected = "aaa[bbbbb[ccc, dd],\n    eee,\n    ffff[gg, hhh, ii]]";
        assert_eq!(expected, pretty(30, doc));
    }
}