Calculate bools while solve(); Add multiple cases while r1cs generation
IfElse only generates flattened code
This commit is contained in:
parent
92e3ede5c0
commit
372d23729d
4 changed files with 122 additions and 63 deletions
|
@ -1,4 +1,6 @@
|
|||
// simple ifelse example
|
||||
// e.g.: x = 5
|
||||
def qeval(x):
|
||||
y = x < 3 ? 1 : 5
|
||||
return x + y
|
||||
z = y < x ? x**3 : y**3
|
||||
return x + y + z
|
||||
|
|
31
src/ast.rs
31
src/ast.rs
|
@ -17,11 +17,11 @@ impl Prog {
|
|||
for def in &self.defs {
|
||||
match *def {
|
||||
Definition::Return(ref expr) => {
|
||||
let s = expr.solve(&witness);
|
||||
let s = expr.solve(&mut witness);
|
||||
witness.insert("~out".to_string(), s);
|
||||
},
|
||||
Definition::Definition(ref id, ref expr) => {
|
||||
let s = expr.solve(&witness);
|
||||
let s = expr.solve(&mut witness);
|
||||
witness.insert(id.to_string(), s);
|
||||
},
|
||||
}
|
||||
|
@ -72,10 +72,31 @@ pub enum Expression {
|
|||
IfElse(Box<Condition>, Box<Expression>, Box<Expression>),
|
||||
}
|
||||
impl Expression {
|
||||
fn solve(&self, inputs: &HashMap<String, i32>) -> i32 {
|
||||
fn solve(&self, inputs: &mut HashMap<String, i32>) -> i32 {
|
||||
match *self {
|
||||
Expression::NumberLiteral(x) => x,
|
||||
Expression::VariableReference(ref var) => inputs[var],
|
||||
Expression::VariableReference(ref var) => {
|
||||
if let None = inputs.get(var) {
|
||||
if var.contains("_b") {
|
||||
let var_name = var.split("_b").collect::<Vec<_>>()[0];
|
||||
let mut num = inputs[var_name];
|
||||
assert!(num >= 0, format!("num < 0: {}", num));
|
||||
// TODO support negative numbers with two's complement!?
|
||||
for i in (0..8).rev() {
|
||||
if 2i32.pow(i) <= num {
|
||||
num -= 2i32.pow(i);
|
||||
inputs.insert(format!("{}_b{}", &var_name, i), 1);
|
||||
} else {
|
||||
inputs.insert(format!("{}_b{}", &var_name, i), 0);
|
||||
}
|
||||
}
|
||||
assert_eq!(num, 0);
|
||||
} else {
|
||||
panic!("Variable not found in inputs: {}", var);
|
||||
}
|
||||
}
|
||||
inputs[var]
|
||||
},
|
||||
Expression::Add(ref x, ref y) => x.solve(inputs) + y.solve(inputs),
|
||||
Expression::Sub(ref x, ref y) => x.solve(inputs) - y.solve(inputs),
|
||||
Expression::Mult(ref x, ref y) => x.solve(inputs) * y.solve(inputs),
|
||||
|
@ -135,7 +156,7 @@ pub enum Condition {
|
|||
Lt(Expression, Expression),
|
||||
}
|
||||
impl Condition {
|
||||
fn solve(&self, inputs: &HashMap<String, i32>) -> bool {
|
||||
fn solve(&self, inputs: &mut HashMap<String, i32>) -> bool {
|
||||
match *self {
|
||||
Condition::Lt(ref lhs, ref rhs) => lhs.solve(inputs) < rhs.solve(inputs),
|
||||
}
|
||||
|
|
113
src/parser.rs
113
src/parser.rs
|
@ -118,6 +118,57 @@ pub fn parse_program(file: File) -> Prog {
|
|||
Prog { id: id, args: args, defs: defs }
|
||||
}
|
||||
|
||||
fn flatten_condition(defs_flattened: &mut Vec<Definition>, num_variables: &mut i32, condition: Condition) -> Expression {
|
||||
match condition {
|
||||
Condition::Lt(lhs, rhs) => {
|
||||
let lhs_flattened = flatten_expression(defs_flattened, num_variables, lhs);
|
||||
let rhs_flattened = flatten_expression(defs_flattened, num_variables, rhs);
|
||||
|
||||
let lhs_name = format!("sym_{}", num_variables);
|
||||
*num_variables += 1;
|
||||
defs_flattened.push(Definition::Definition(lhs_name.to_string(), lhs_flattened));
|
||||
let rhs_name = format!("sym_{}", num_variables);
|
||||
*num_variables += 1;
|
||||
defs_flattened.push(Definition::Definition(rhs_name.to_string(), rhs_flattened));
|
||||
|
||||
let cond_result = format!("sym_{}", num_variables);
|
||||
*num_variables += 1;
|
||||
defs_flattened.push(Definition::Definition(
|
||||
cond_result.to_string(),
|
||||
Sub(
|
||||
box VariableReference(lhs_name.to_string()),
|
||||
box VariableReference(rhs_name.to_string())
|
||||
)
|
||||
));
|
||||
let bits = 8;
|
||||
for i in 0..bits {
|
||||
let new_name = format!("{}_b{}", &cond_result, i);
|
||||
defs_flattened.push(Definition::Definition(
|
||||
new_name.to_string(),
|
||||
Mult(
|
||||
box VariableReference(new_name.to_string()),
|
||||
box VariableReference(new_name.to_string())
|
||||
)
|
||||
));
|
||||
}
|
||||
let mut expr = VariableReference(format!("{}_b0", &cond_result)); // * 2^0
|
||||
for i in 1..bits {
|
||||
expr = Add(
|
||||
box Mult(
|
||||
box VariableReference(format!("{}_b{}", &cond_result, i)),
|
||||
box NumberLiteral(2i32.pow(i))
|
||||
),
|
||||
box expr
|
||||
);
|
||||
}
|
||||
defs_flattened.push(Definition::Definition(cond_result.to_string(), expr));
|
||||
|
||||
let cond_true = format!("{}_b{}", &cond_result, bits - 1);
|
||||
VariableReference(cond_true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn flatten_expression(defs_flattened: &mut Vec<Definition>, num_variables: &mut i32, expr: Expression) -> Expression {
|
||||
match expr {
|
||||
x @ NumberLiteral(_) |
|
||||
|
@ -249,74 +300,24 @@ fn flatten_expression(defs_flattened: &mut Vec<Definition>, num_variables: &mut
|
|||
},
|
||||
IfElse(box condition, consequent, alternative) => {
|
||||
let condition_true = flatten_condition(defs_flattened, num_variables, condition);
|
||||
let new_name = format!("sym_{}", num_variables);
|
||||
*num_variables += 1;
|
||||
// condition_false = 1 - condition_true
|
||||
defs_flattened.push(Definition::Definition(new_name.to_string(), Sub(box NumberLiteral(1), box condition_true.clone())));
|
||||
let condition_false = VariableReference(new_name);
|
||||
// (condition_true * consequent) + (condition_false * alternatuve)
|
||||
flatten_expression(
|
||||
defs_flattened,
|
||||
num_variables,
|
||||
Add(
|
||||
box Mult(box condition_true.clone(), consequent),
|
||||
box Mult(
|
||||
box Sub(box NumberLiteral(1), box condition_true),
|
||||
alternative
|
||||
)
|
||||
box Mult(box condition_true, consequent),
|
||||
box Mult(box condition_false, alternative)
|
||||
)
|
||||
)
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn flatten_condition(defs_flattened: &mut Vec<Definition>, num_variables: &mut i32, condition: Condition) -> Expression {
|
||||
match condition {
|
||||
Condition::Lt(lhs, rhs) => {
|
||||
let lhs_flattened = flatten_expression(defs_flattened, num_variables, lhs);
|
||||
let rhs_flattened = flatten_expression(defs_flattened, num_variables, rhs);
|
||||
|
||||
let lhs_name = format!("sym_{}", num_variables);
|
||||
*num_variables += 1;
|
||||
defs_flattened.push(Definition::Definition(lhs_name.to_string(), lhs_flattened));
|
||||
let rhs_name = format!("sym_{}", num_variables);
|
||||
*num_variables += 1;
|
||||
defs_flattened.push(Definition::Definition(rhs_name.to_string(), rhs_flattened));
|
||||
|
||||
let cond_result = format!("sym_{}", num_variables);
|
||||
*num_variables += 1;
|
||||
defs_flattened.push(Definition::Definition(
|
||||
cond_result.to_string(),
|
||||
Sub(
|
||||
box VariableReference(lhs_name.to_string()),
|
||||
box VariableReference(rhs_name.to_string())
|
||||
)
|
||||
));
|
||||
let bits = 8;
|
||||
for i in 0..bits {
|
||||
let new_name = format!("{}_b{}", &cond_result, i);
|
||||
defs_flattened.push(Definition::Definition(
|
||||
new_name.to_string(),
|
||||
Mult(
|
||||
box VariableReference(new_name.to_string()),
|
||||
box VariableReference(new_name.to_string())
|
||||
)
|
||||
));
|
||||
}
|
||||
let mut expr = VariableReference(format!("{}_b0", &cond_result)); // * 2^0
|
||||
for i in 1..bits {
|
||||
expr = Add(
|
||||
box Mult(
|
||||
box VariableReference(format!("{}_b{}", &cond_result, i)),
|
||||
box NumberLiteral(2i32.pow(i))
|
||||
),
|
||||
box expr
|
||||
);
|
||||
}
|
||||
defs_flattened.push(Definition::Definition(cond_result.to_string(), expr));
|
||||
|
||||
let cond_true = format!("{}_b{}", &cond_result, bits - 1);
|
||||
VariableReference(cond_true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn flatten_program(prog: Prog) -> Prog {
|
||||
let mut defs_flattened = Vec::new();
|
||||
let mut num_variables: i32 = 0;
|
||||
|
|
37
src/r1cs.rs
37
src/r1cs.rs
|
@ -65,6 +65,39 @@ fn count_variables_add(expr: Expression) -> HashMap<String, i32> {
|
|||
let var = count.entry(v).or_insert(0);
|
||||
*var += n;
|
||||
},
|
||||
(VariableReference(v1), Mult(box NumberLiteral(n), box VariableReference(v2))) |
|
||||
(VariableReference(v1), Mult(box VariableReference(v2), box NumberLiteral(n))) |
|
||||
(Mult(box NumberLiteral(n), box VariableReference(v2)), VariableReference(v1)) |
|
||||
(Mult(box VariableReference(v2), box NumberLiteral(n)), VariableReference(v1)) => {
|
||||
{
|
||||
let var = count.entry(v1).or_insert(0);
|
||||
*var += 1;
|
||||
}
|
||||
let var = count.entry(v2).or_insert(0);
|
||||
*var += n;
|
||||
},
|
||||
(e @ Add(..), Mult(box NumberLiteral(n), box VariableReference(v))) |
|
||||
(e @ Add(..), Mult(box VariableReference(v), box NumberLiteral(n))) |
|
||||
(Mult(box NumberLiteral(n), box VariableReference(v)), e @ Add(..)) |
|
||||
(Mult(box VariableReference(v), box NumberLiteral(n)), e @ Add(..)) => {
|
||||
{
|
||||
let var = count.entry(v).or_insert(0);
|
||||
*var += n;
|
||||
}
|
||||
let vars = count_variables_add(e);
|
||||
for (key, value) in &vars {
|
||||
let val = count.entry(key.to_string()).or_insert(0);
|
||||
*val += *value;
|
||||
}
|
||||
},
|
||||
(Mult(box VariableReference(v1), box NumberLiteral(n1)), Mult(box VariableReference(v2), box NumberLiteral(n2))) => {
|
||||
{
|
||||
let var = count.entry(v1).or_insert(0);
|
||||
*var += n1;
|
||||
}
|
||||
let var = count.entry(v2).or_insert(0);
|
||||
*var += n2;
|
||||
},
|
||||
e @ _ => panic!("Error: Add({}, {})", e.0, e.1),
|
||||
}
|
||||
},
|
||||
|
@ -80,7 +113,9 @@ fn swap_sub(lhs: &Expression, rhs: &Expression) -> (Expression, Expression) {
|
|||
(v1 @ NumberLiteral(_), v2 @ NumberLiteral(_)) |
|
||||
(v1 @ VariableReference(_), v2 @ NumberLiteral(_)) |
|
||||
(v1 @ NumberLiteral(_), v2 @ VariableReference(_)) |
|
||||
(v1 @ VariableReference(_), v2 @ VariableReference(_)) => (v1, v2),
|
||||
(v1 @ VariableReference(_), v2 @ VariableReference(_)) |
|
||||
(v1 @ VariableReference(_), v2 @ Mult(..)) |
|
||||
(v1 @ Mult(..), v2 @ VariableReference(_)) => (v1, v2),
|
||||
(var @ VariableReference(_), Add(left, box right)) => {
|
||||
let (l, r) = swap_sub(&var, &right);
|
||||
(l, Add(left, box r))
|
||||
|
|
Loading…
Reference in a new issue