main.rs

extern crate core;

use std::fmt::Display;
use std::fmt::Formatter;
use core::fmt;
use core::mem;

const TCBLIKE: &'static str = "
\\set aid random(1, 100000 * $scale)
\\set bid random(1, 1 * $scale)
\\set tid random(1, 10 * $scale)
\\set delta random(-5000, 5000)
BEGIN;
UPDATE pgbench_accounts SET abalance = abalance + $delta WHERE aid = :aid;
SELECT abalance FROM pgbench_accounts WHERE aid = $aid;
UPDATE pgbench_tellers SET tbalance = tbalance + $delta WHERE tid = :tid;
UPDATE pgbench_branches SET bbalance = bbalance + $delta WHERE bid = :bid;
INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES ($tid, $bid, $aid, $delta, CURRENT_TIMESTAMP);
END;";

struct MetaCommandSet {
    var: usize,
//    expression: Expression,
}

struct Variables {
    vars: Vec<String>
}
impl Variables {
    fn get_or_create(&mut self, name: String) -> usize {
        let mut index = 0;
        while index < self.vars.len() {
            if self.vars[index] == name {
                return index
            }
            index += 1;
        }
        self.vars.push(name);
        return self.vars.len()
    }
}

trait Command {

}

#[derive(Debug)]
enum ParserState {
    Idle,
    MetaCommandPre,
    MetaCommandSetVar,
    MetaCommandExpression,
    CypherCommand,
}

fn main() {
    let script = TCBLIKE;
    let chars = script.chars();

    let mut vars = Variables {
        vars: Vec::new(),
    };

    let mut mark = 0;
    let mut cursor = 0;
    let mut state = ParserState::Idle;

    let mut meta_cmd_set: MetaCommandSet;

    for c in chars {
        state = match state {
            ParserState::Idle => match c {
                '\\' => {
                    mark = cursor + 1;
                    ParserState::MetaCommandPre
                }
                ' ' | '\n' => {
                    mark = cursor;
                    ParserState::Idle
                }
                _ => {
                    mark = cursor;
                    ParserState::CypherCommand
                }
            }
            ParserState::MetaCommandPre => match c {
                ' ' => {
                    let cmd: String = script.chars().skip(mark).take(cursor - mark).collect();
                    match cmd.as_ref() {
                        "set" => {
                            mark = cursor + 1;
                            ParserState::MetaCommandSetVar
                        }
                        _ => panic!("Unknown meta command {}", cmd)
                    }
                }
                '\n' => {
                    panic!("Expected meta command, like '\\set somevar ..', to follow \\, but got newline")
                }
                _ => ParserState::MetaCommandPre
            }
            ParserState::MetaCommandSetVar => match c {
                ' ' => {
                    let var_name: String = script.chars().skip(mark).take(cursor - mark).collect();
                    mark = cursor + 1;
                    meta_cmd_set = MetaCommandSet{
                        var: vars.get_or_create(var_name),
                    };
                    ParserState::MetaCommandExpression
                }
                '\n' => {
                    panic!("Expected variable name, like '\\set somevar ..', to follow \\set .., but got newline")
                }
                _ => ParserState::MetaCommandSetVar
            }
            ParserState::MetaCommandExpression => match c {
                '\n' => {
                    let expression_str: String = script.chars().skip(mark).take(cursor - mark).collect();
                    let expression = parse_expression(expression_str);

                    print!("{}", expression.describe());

                    ParserState::Idle
                }
                _ => ParserState::MetaCommandExpression
            }
            ParserState::CypherCommand => match c {
                ';' => {
                    let command: &str = unsafe { script.slice_unchecked(mark, cursor) };
                    println!("{}", command);
                    ParserState::Idle
                }
                _ => ParserState::CypherCommand
            }
        };
//        println!("{} {} - {:?}", cursor, c, state);
        cursor += 1;
    };
}

trait Expression {
    fn describe(&self) -> String;
}

enum FunctionName {
    Random,
    RandomGaussian,
}

impl FunctionName {
    pub fn from(name: String) -> FunctionName {
        match name.as_ref() {
            "random" => FunctionName::Random,
            "random_gaussian" => FunctionName::RandomGaussian,
            _ => panic!("Unknown function {}", name)
        }
    }
}

impl Display for FunctionName {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        match self {
            FunctionName::Random => write!(f, "random"),
            FunctionName::RandomGaussian => write!(f, "random_gaussian")
        }
    }
}

struct ExprFuncCall {
    function: FunctionName,
    args: Vec<Box<Expression>>,
}

impl ExprFuncCall {
    pub fn new(function: FunctionName) -> Box<ExprFuncCall> {
        Box::new(ExprFuncCall{
            function: function,
            args: Vec::new(),
        })
    }
}

impl Expression for ExprFuncCall {
    fn describe(&self) -> String {
        let mut args = String::from("");
        let mut first = true;
        for arg in &self.args {
            if !first {
                args += ", ";
            }
            args += arg.describe().as_ref();
            first = false;
        }
        format!("{}({})", self.function.to_string(), args)
    }
}

struct ExprVarLiteral {
    varName: String,
}

impl Expression for ExprVarLiteral {
    fn describe(&self) -> String {
        format!("${}", self.varName)
    }
}

struct ExprIntLiteral {
    val: i64,
}

impl ExprIntLiteral {
    pub fn from(raw: String) -> ExprIntLiteral {
        ExprIntLiteral{
            val: raw.parse::<i64>().expect(&raw),
        }
    }
}

impl Expression for ExprIntLiteral {
    fn describe(&self) -> String {
        format!("{}", self.val)
    }
}

#[derive(Debug)]
enum ExpressionParserState {
    Idle,
    Int,
    Var,
}

struct ExpressionParserCtx {
    raw: String,
    mark: usize,
    cursor: usize,
    stack: Vec<Box<ExprFuncCall>>,
    current: Option<Box<ExprFuncCall>>,
}

fn parse_expression(raw: String) -> Box<Expression> {
    fn end_function(ctx: &mut ExpressionParserCtx) -> ExpressionParserState {
        // At the end of a function, check if there is an outer function call
        // and, if so, our function should be an argument to that outer function,
        // and the outer function becomes the current function we're working on..
        match ctx.stack.pop() {
            Some(mut e) => {
                // Yep, add the current function as an argument, and make the parent
                // into the current function.
                e.args.push(ctx.current.take().unwrap());
                ctx.current = Some(e)
            }
            None => ()
        }
        ExpressionParserState::Idle
    }

    fn end_var(ctx: &mut ExpressionParserCtx) -> Box<Expression> {
        let expr = Box::new(ExprVarLiteral{varName: token(ctx)});
        ctx.mark = ctx.cursor + 1;
        expr
    }

    fn end_var_arg(ctx: &mut ExpressionParserCtx) {
        let expr = end_var(ctx);
        ctx.current.as_mut().unwrap().args.push(expr);
    }

    fn end_int(ctx: &mut ExpressionParserCtx) -> Box<Expression> {
        let expr = Box::new(ExprIntLiteral::from(token(ctx)));
        ctx.mark = ctx.cursor + 1;
        expr
    }

    fn end_int_arg(ctx: &mut ExpressionParserCtx) {
        let expr = end_int(ctx);
        ctx.current.as_mut().unwrap().args.push(expr);
    }

    fn token(ctx: &ExpressionParserCtx) -> String {
        ctx.raw.chars().skip(ctx.mark).take(ctx.cursor - ctx.mark).collect()
    }

    let mut chars = raw.chars();
    let mut state = ExpressionParserState::Idle;

    let mut ctx = ExpressionParserCtx{
        raw: raw.clone(),
        mark: 0,
        cursor: 0,
        stack: Vec::new(),
        current: None,
    };

    let mut c;
    loop {
        c = chars.next().unwrap_or('\0');
        state = match state {
            ExpressionParserState::Idle => match c {
                '(' => {
                    let func_name: String = token(&ctx);
                    ctx.mark = ctx.cursor + 1;
                    let parent = mem::replace(&mut ctx.current, Some(ExprFuncCall::new( FunctionName::from(func_name) )));

                    match parent {
                        Some(f) => {
                            ctx.stack.push(f)
                        }
                        _ => ()
                    }

                    ExpressionParserState::Idle
                },
                ')' => end_function(&mut ctx),
                '0'...'9' | '-' => ExpressionParserState::Int,
                ' ' => {
                    ctx.mark = ctx.cursor + 1;
                    ExpressionParserState::Idle
                }
                '$' => {
                    ctx.mark = ctx.cursor + 1;
                    ExpressionParserState::Var
                }
                _ => ExpressionParserState::Idle,
            }
            ExpressionParserState::Int => match c {
                '0'...'9' => ExpressionParserState::Int,
                ')' => {
                    end_int_arg(&mut ctx);
                    end_function(&mut ctx)
                }
                ',' => {
                    end_int_arg(&mut ctx);
                    ExpressionParserState::Idle
                }
                '\0' => return end_int(&mut ctx),
                _ => panic!("Don't know about {} in {} from {:?}", c, raw, state)
            },
            ExpressionParserState::Var => match c {
                '0'...'z' => ExpressionParserState::Var,
                ')' => {
                    end_var_arg(&mut ctx);
                    end_function(&mut ctx)
                }
                ',' => {
                    end_var_arg(&mut ctx);
                    ExpressionParserState::Idle
                }
                '\0' => return end_var(&mut ctx),
                _ => panic!("Don't know about '{}' in {} from {:?}", c, raw, state)
            },
        };

        if c == '\0' {
            break;
        }

        ctx.cursor += 1;
    }

    return ctx.current.expect("Should be defined")
}

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

    #[test]
    fn parseInt() {
        assert_eq!(parse_expression(String::from("1")).describe(), "1");
        assert_eq!(parse_expression(String::from("0")).describe(), "0");
        assert_eq!(parse_expression(String::from("-1")).describe(), "-1");
        assert_eq!(parse_expression(String::from("234")).describe(), "234");
    }

    #[test]
    fn parseVar() {
        assert_eq!(parse_expression(String::from("$banana")).describe(), "$banana");
    }

    #[test]
    fn parseNoArgFunc() {
        assert_eq!(parse_expression(String::from("random()")).describe(), "random()");
    }

    #[test]
    fn parseFuncWithVariableArg() {
        assert_eq!(parse_expression(String::from("random($banana)")).describe(), "random($banana)");
        assert_eq!(parse_expression(String::from("random($boo, $banana)")).describe(), "random($boo, $banana)");
        assert_eq!(parse_expression(String::from("random(11, $banana)")).describe(), "random(11, $banana)");
        assert_eq!(parse_expression(String::from("random($banana, 11)")).describe(), "random($banana, 11)");
    }

    #[test]
    fn parseFuncWithFuncArg() {
        assert_eq!(parse_expression(String::from("random(random())")).describe(), "random(random())");
    }

    #[test]
    fn parseFuncWithIntArg() {
        assert_eq!(parse_expression(String::from("random(1234)")).describe(), "random(1234)");
        assert_eq!(parse_expression(String::from("random(0)")).describe(), "random(0)");
        assert_eq!(parse_expression(String::from("random(-1)")).describe(), "random(-1)");


        assert_eq!(parse_expression(String::from("random(1234, 4321)")).describe(), "random(1234, 4321)");
    }

    #[test]
    fn parseFuncWithNestedFuncIntArg() {
        assert_eq!(parse_expression(String::from("random_gaussian(random(1234))")).describe(), "random_gaussian(random(1234))");

        assert_eq!(parse_expression(String::from("random_gaussian(random(1234), 1)")).describe(), "random_gaussian(random(1234), 1)");
    }
}