test.rs

#![feature(plugin_registrar, rustc_private, quote)]

extern crate aster;
extern crate petgraph;
extern crate rustc_plugin;
extern crate syntax;

use std::mem;

use syntax::ast;
use syntax::codemap::Span;
use syntax::ext::base::{
    Annotatable,
    ExtCtxt,
    MultiModifier,
};
use syntax::ext::build::AstBuilder;
use syntax::ptr::P;

use rustc_plugin::Registry;
use petgraph::graph::{Graph, NodeIndex};

/*
pub fn expand_stateful<'cx>(
    cx: &'cx mut ExtCtxt,
    sp: Span,
    _: &[ast::TokenTree],
) -> Box<MacResult + 'cx> {
    let expr = cx.expr_str(sp, InternedString::new("hello world"));
    MacEager::expr(expr)
}
*/

struct CFGBuilder {
    graph: Graph<Node, ()>,
    entry: NodeIndex,
    exit: NodeIndex,
}

impl CFGBuilder {
    fn new() -> Self {
        let mut graph = Graph::new();
        let entry = graph.add_node(Node::Entry);
        let exit = graph.add_node(Node::Exit);
        CFGBuilder {
            graph: graph,
            entry: entry,
            exit: exit,
        }
    }
}

struct CFG {
    graph: Graph<Node, ()>,
    entry: NodeIndex,
    exit: NodeIndex,
}

enum Node {
    BasicBlock(BasicBlock),
    Entry,
    Exit,
}

struct BasicBlock {
    stmts: Vec<P<ast::Stmt>>,
    expr: Option<P<ast::Expr>>,
}

impl BasicBlock {
    fn new() -> Self {
        BasicBlock {
            stmts: Vec::new(),
            expr: None,
        }
    }

    fn block(&mut self, block: &ast::Block, pred: NodeIndex) -> NodeIndex {
        let mut stmts_exit = pred;
        for stmt in &block.stmts {
            stmts_exit = self.stmt(stmt, stmts_exit);
        }

        let expr_exit = self.opt_expr(&blk.expr, stmts_exit);
        expr_exit
    }

    fn stmt(&mut self, stmt: &P<ast::Stmt>, pred: NodeIndex, bb: BasicBlock) -> (NodeIndex,  {
        match stmt.node {
            ast::Stmt_::StmtSemi(ref expr, _) => {
                let exit = self.expr(expr, pred);
            }
            _ => {
                panic!()
            }
        }
    }

    fn expr(&mut self, expr: &P<ast::Expr>, pred: NodeIndex) -> NodeIndex {
        let exit = match expr.node {
            ast::Expr_::ExprBlock(ref block) => {
                self.block(block, pred)
            }
            ast::Expr_::ExprRet(ref e) => {
            }
            _ => {
                
            }
        };

        self.graph.add_edge(
    }

    fn get_node(&mut self, index: NodeIndex) -> &mut 

    fn add_node(&mut self, Node::BasicBlock(
}

/*
fn visit_stmt(graph: &mut Graph<BasicBlock, ()>,
       parent_index: NodeIndex,
       stmt: P<ast::Stmt>) -> NodeIndex {
    let builder = aster::AstBuilder::new();

    match stmt.node {
        ast::Stmt_::StmtExpr(ref expr, _) | ast::Stmt_::StmtSemi(ref expr, _) => {
            match expr.node {
                ast::Expr_::ExprRet(ref expr) => {
                    let node = BasicBlock::new();
                    let node_index = graph.add_node(node);
                    graph.add_edge(parent_index, node_index, ());


                    let parent = graph.node_weight_mut(parent_index).unwrap();
                    parent.expr = match *expr {
                        Some(ref expr) => Some(expr.clone()),
                        None => Some(builder.expr().unit()),
                    };

                    return node_index;
                }
                _ => { }
            }
        }
        _ => { }
    }

    let parent = graph.node_weight_mut(parent_index).unwrap();
    parent.stmts.push(stmt.clone());
    parent_index
}

fn visit_stmts(cx: &ExtCtxt,
               block: &ast::Block) -> Graph<BasicBlock, ()>, NodeIndex) {
    let mut graph = Graph::new();

    let entry = BasicBlock::new();
    let entry_index = graph.add_node(entry);

    let mut exit = BasicBlock::new();


    let mut index = entry_index;

    for stmt in block.stmts.iter() {
        index = visit_stmt(graph, entry_index, stmt.clone());
    }

    match block.expr {
        Some(ref expr) => {
            let bb = BasicBlock::new();
            exit.expr = Some(expr.clone());
        }
        None => { }
    }

    (graph, entry_index)
}
*/


fn expand_state_machine(cx: &mut ExtCtxt,
                        _sp: Span,
                        meta_item: &ast::MetaItem,
                        annotatable: Annotatable) -> Annotatable {
    let builder = aster::AstBuilder::new();

    let item = match annotatable {
        Annotatable::Item(item) => item,
        _ => {
            cx.span_err(
                meta_item.span,
                "`state_machine` may only be applied to functions");

            return annotatable;
        }
    };

    let ident = item.ident;

    let (fn_decl, block) = match item.node {
        ast::ItemFn(ref fn_decl, _, _, _, _, ref block) => (fn_decl, block),
        _ => {
            cx.span_err(
                meta_item.span,
                "`state_machine` may only be applied to functions");

            return Annotatable::Item(item.clone());
        }
    };

    let ret_ty = match fn_decl.output {
        ast::FunctionRetTy::NoReturn(..) => {
            cx.span_err(
                meta_item.span,
                "`state_machine` cannot return `!` types");

            return Annotatable::Item(item.clone());
        }
        ast::FunctionRetTy::DefaultReturn(span) => builder.ty().span(span).unit(),
        ast::FunctionRetTy::Return(ref ty) => ty.clone(),
    };

    /*
    let mut graph = Graph::<BasicBlock, ()>::new();
    bar(&mut graph, &block.stmts[..]);
    */

    let cfg_builder = CFGBuilder::new();


    let mut state_blocks = vec![
        quote_block!(cx, {
            return (::std::option::Option::None::<$ret_ty>, State::Exit)
        }),
    ];

    let mut current_stmts = vec![];

    for stmt in block.stmts.iter().rev() {
        match stmt.node {
            ast::Stmt_::StmtExpr(ref e, _) | ast::Stmt_::StmtSemi(ref e, _) => {
                match e.node {
                    ast::Expr_::ExprRet(ref e) => {
                        let src_node = state_blocks.len();
                        let dst_node = state_blocks.len() + 1;

                        let target_state = state_blocks.len();
                        let target_state_id = builder.id(format!("State{}", target_state));

                        let mut stmts = mem::replace(&mut current_stmts, vec![]);
                        stmts.reverse();

                        state_blocks.push(quote_block!(cx, {
                            $stmts
                            return (::std::option::Option::Some($e), State::$target_state_id);
                        }));
                    }
                    _ => {
                        current_stmts.push(stmt.clone());
                    }
                }
            }
            _ => {
                current_stmts.push(stmt.clone());
            }
        }
    }

    if !current_stmts.is_empty() {
        current_stmts.reverse();

        state_blocks.push(quote_block!(cx, {
            $current_stmts
        }));
    }

    let state_ids = (0 .. state_blocks.len())
        .map(|i| {
            builder.id(format!("State{}", i))
        })
        .collect::<Vec<_>>();

    let state_enum = builder.item().enum_("State")
        .ids(state_ids.iter())
        .id("Exit")
        .build();

    let mut state_arms = state_ids.iter()
        .zip(state_blocks.iter().rev())
        .map(|(state, block)| quote_arm!(cx, State::$state => $block))
        .collect::<Vec<_>>();

    state_arms.push(quote_arm!(cx,
        State::Exit => {
            return (::std::option::Option::None::<$ret_ty>, State::Exit)
        }
    ));

    let item = quote_item!(cx,
        fn $ident() -> ::std::boxed::Box<::std::iter::Iterator<Item=$ret_ty>> {
            struct Wrapper<S, F> {
                state: S,
                next: F,
            }

            impl<S, T, F> Wrapper<S, F>
                where F: Fn(S) -> (Option<T>, S),
            {
                fn new(initial_state: S, next: F) -> Self {
                    Wrapper {
                        state: initial_state,
                        next: next,
                    }
                }
            }

            impl<S, T, F> Iterator for Wrapper<S, F>
                where S: Default,
                      F: Fn(S) -> (Option<T>, S)
            {
                type Item = T;

                fn next(&mut self) -> Option<Self::Item> {
                    let old_state = ::std::mem::replace(&mut self.state, S::default());
                    let (value, next_state) = (self.next)(old_state);
                    self.state = next_state;
                    value
                }
            }

            $state_enum

            impl Default for State {
                fn default() -> Self {
                    State::Exit
                }
            }

            Box::new(Wrapper::new(
                State::Exit,
                |mut state| {
                    loop {
                        match state {
                            $state_arms
                        }
                    }
                }
            ))
        }
    ).unwrap();

    Annotatable::Item(item)

    //let expr = cx.expr_str(sp, InternedString::new("hello world"));
    //push(Annotatable(
}

#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(registry: &mut rustc_plugin::Registry) {
    let builder = aster::AstBuilder::new();

    registry.register_syntax_extension(builder.name("state_machine"),
                                       MultiModifier(Box::new(expand_state_machine)));
}