Skip to content

Commit 6105a83

Browse files
windelbouwmanwindelbouwman
committed
Initial variant of list comprehensions
1 parent adfb750 commit 6105a83

8 files changed

Lines changed: 234 additions & 5 deletions

File tree

‎parser/src/ast.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,16 @@ pub enum Expression {
165165
Dict {
166166
elements: Vec<(Expression, Expression)>,
167167
},
168+
ListComprehension {
169+
element: Box<Expression>,
170+
generators: Vec<Comprehension>,
171+
},
172+
/*
173+
DictComprehension {
174+
key: Box<Expression>,
175+
value: Box<Expression>,
176+
generators: Vec<Comprehension>,
177+
},*/
168178
Slice {
169179
elements: Vec<Expression>,
170180
},
@@ -188,6 +198,13 @@ pub enum Expression {
188198
None,
189199
}
190200

201+
#[derive(Debug, PartialEq, Clone)]
202+
pub struct Comprehension {
203+
pub target: Vec<Expression>, // TODO: should this be a non-vector?
204+
pub iter: Expression,
205+
pub ifs: Vec<Expression>,
206+
}
207+
191208
#[derive(Debug, PartialEq, Clone)]
192209
pub struct Keyword {
193210
pub name: Option<String>,

‎parser/src/parser.rs‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ pub fn parse_expression(source: &str) -> Result<ast::Expression, String> {
7171
#[cfg(test)]
7272
mod tests {
7373
use super::ast;
74+
use super::parse_expression;
7475
use super::parse_program;
7576
use super::parse_statement;
7677

@@ -324,4 +325,27 @@ mod tests {
324325
})
325326
)
326327
}
328+
329+
#[test]
330+
fn test_parse_list_comprehension() {
331+
let source = String::from("[x for y in z]");
332+
let parse_ast = parse_expression(&source).unwrap();
333+
assert_eq!(
334+
parse_ast,
335+
ast::Expression::ListComprehension {
336+
element: Box::new(ast::Expression::Identifier {
337+
name: "x".to_string()
338+
}),
339+
generators: vec![ast::Comprehension {
340+
target: vec![ast::Expression::Identifier {
341+
name: "y".to_string()
342+
},],
343+
iter: ast::Expression::Identifier {
344+
name: "z".to_string()
345+
},
346+
ifs: vec![],
347+
}],
348+
}
349+
);
350+
}
327351
}

‎parser/src/python.lalrpop‎

Lines changed: 66 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -567,11 +567,13 @@ Atom: ast::Expression = {
567567
<s:String> => ast::Expression::String { value: s },
568568
<n:Number> => ast::Expression::Number { value: n },
569569
<i:Identifier> => ast::Expression::Identifier { name: i },
570-
"[" <e:TestList?> <_trailing_comma:","?> "]" => {
571-
match e {
572-
None => ast::Expression::List { elements: Vec::new() },
573-
Some(elements) => ast::Expression::List { elements },
574-
}
570+
"[" <e:TestListComp?> "]" => {
571+
let elements = e.unwrap_or(Vec::new());
572+
ast::Expression::List { elements }
573+
},
574+
"[" <e:TestListComp2> "]" => {
575+
// List comprehension:
576+
e
575577
},
576578
"(" <e:TestList?> <trailing_comma:","?> ")" => {
577579
match e {
@@ -592,6 +594,22 @@ Atom: ast::Expression = {
592594
"None" => ast::Expression::None,
593595
};
594596

597+
TestListComp: Vec<ast::Expression> = {
598+
<e:TestList> <_trailing_comma:","?> => {
599+
e
600+
},
601+
};
602+
603+
TestListComp2: ast::Expression = {
604+
<e:Test> <c:CompFor> => {
605+
// vec![e]
606+
ast::Expression::ListComprehension {
607+
element: Box::new(e),
608+
generators: c,
609+
}
610+
},
611+
};
612+
595613
TestDict: Vec<(ast::Expression, ast::Expression)> = {
596614
<e1:DictEntry> <e2:("," DictEntry)*> <_trailing_comma:","?> => {
597615
let mut d = vec![e1];
@@ -617,6 +635,49 @@ TestList: Vec<ast::Expression> = {
617635
}
618636
};
619637

638+
// Comprehensions:
639+
CompIter: (Option<Vec<ast::Comprehension>>, Option<ast::Expression>) = {
640+
// CompIf,
641+
<c:CompFor> => (Some(c), None),
642+
};
643+
644+
CompFor: Vec<ast::Comprehension> = {
645+
"for" <e:ExpressionList> "in" <i:OrTest> <c2:CompIter?> => {
646+
match c2 {
647+
None => {
648+
vec![ast::Comprehension {
649+
target: e,
650+
iter: i,
651+
ifs: vec![],
652+
}]
653+
},
654+
Some((Some(mut cphs), None)) => {
655+
let mut res = vec![];
656+
res.append(&mut cphs);
657+
res.push(ast::Comprehension {
658+
target: e,
659+
iter: i,
660+
ifs: vec![],
661+
});
662+
res
663+
},
664+
_ => {
665+
panic!("TODO");
666+
},
667+
}
668+
}
669+
};
670+
671+
ExpressionNoCond: ast::Expression = {
672+
OrTest,
673+
};
674+
675+
//CompIf: ast::Expression = {
676+
// "if" <c:ExpressionNoCond> <c2:CompIter?> => {
677+
// c
678+
// }
679+
//};
680+
620681
ArgumentList: (Vec<ast::Expression>, Vec<ast::Keyword>) = {
621682
<e: Comma<FunctionArgument>> => {
622683
let mut args = vec![];

‎tests/snippets/comprehensions.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
2+
x = [1, 2, 3]
3+
4+
y = [a*a+1 for a in x]
5+
assert y == [2, 5, 10]
6+

‎vm/src/bytecode.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,9 @@ pub enum Instruction {
140140
BuildSlice {
141141
size: usize,
142142
},
143+
ListAppend {
144+
i: usize,
145+
},
143146
PrintExpr,
144147
LoadBuildClass,
145148
StoreLocals,

‎vm/src/compile.rs‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -250,6 +250,7 @@ impl Compiler {
250250
self.emit(Instruction::ForIter);
251251

252252
// Start of loop iteration, set targets:
253+
// TODO: can we use compile_store here?
253254
for t in target {
254255
match t {
255256
ast::Expression::Identifier { name } => {
@@ -878,6 +879,107 @@ impl Compiler {
878879
flags: bytecode::FunctionOpArg::empty(),
879880
});
880881
}
882+
ast::Expression::ListComprehension {
883+
element,
884+
generators,
885+
} => {
886+
// Okay, compile a new function which can be called
887+
888+
// For now only support one generator:
889+
if generators.len() != 1 {
890+
unimplemented!("Only 1 level deep list comprehensions are implemented now");
891+
}
892+
let generator = &generators[0];
893+
894+
// Create magnificent function <listcomp>:
895+
self.code_object_stack.push(CodeObject::new(
896+
vec![".0".to_string()],
897+
self.source_path.clone(),
898+
"<listcomp>".to_string(),
899+
));
900+
901+
// Create empty list:
902+
self.emit(Instruction::BuildList { size: 0 });
903+
904+
// Load 'append' method:
905+
//self.emit(Instruction::LoadAttr {
906+
// name: "append".to_string(),
907+
//});
908+
909+
// Load iterator onto stack:
910+
self.emit(Instruction::LoadName {
911+
name: String::from(".0"),
912+
});
913+
914+
// Setup for loop:
915+
let start_label = self.new_label();
916+
let end_label = self.new_label();
917+
self.emit(Instruction::SetupLoop {
918+
start: start_label,
919+
end: end_label,
920+
});
921+
self.set_label(start_label);
922+
self.emit(Instruction::ForIter);
923+
924+
// TODO: can we use compile_store here?
925+
for t in &generator.target {
926+
match t {
927+
ast::Expression::Identifier { name } => {
928+
self.emit(Instruction::StoreName {
929+
name: name.to_string(),
930+
});
931+
}
932+
_ => panic!("Not impl"),
933+
}
934+
}
935+
936+
// Evaluate element:
937+
self.compile_expression(element);
938+
939+
// List append:
940+
self.emit(Instruction::ListAppend { i: 2 });
941+
942+
// Repeat:
943+
self.emit(Instruction::Jump {
944+
target: start_label,
945+
});
946+
947+
// End of for loop:
948+
self.set_label(end_label);
949+
self.emit(Instruction::PopBlock);
950+
951+
// Return freshly filled list:
952+
self.emit(Instruction::ReturnValue);
953+
954+
// Fetch code for listcomp function:
955+
let code = self.code_object_stack.pop().unwrap();
956+
957+
// List comprehension code:
958+
self.emit(Instruction::LoadConst {
959+
value: bytecode::Constant::Code { code: code },
960+
});
961+
962+
// List comprehension function name:
963+
self.emit(Instruction::LoadConst {
964+
value: bytecode::Constant::String {
965+
value: String::from("<listcomp>"),
966+
},
967+
});
968+
969+
// Turn code object into function object:
970+
self.emit(Instruction::MakeFunction {
971+
flags: bytecode::FunctionOpArg::empty(),
972+
});
973+
974+
// Evaluate iterated item:
975+
self.compile_expression(&generator.iter);
976+
977+
// Get iterator / turn item into an iterator
978+
self.emit(Instruction::GetIter);
979+
980+
// Call just created <listcomp> function:
981+
self.emit(Instruction::CallFunction { count: 1 });
982+
}
881983
ast::Expression::IfExpression { test, body, orelse } => {
882984
let no_label = self.new_label();
883985
let end_label = self.new_label();

‎vm/src/frame.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,10 @@ impl Frame {
110110
pub fn last_value(&self) -> PyObjectRef {
111111
self.stack.last().unwrap().clone()
112112
}
113+
114+
pub fn nth_value(&self, depth: usize) -> PyObjectRef {
115+
self.stack[self.stack.len() - depth - 1].clone()
116+
}
113117
}
114118

115119
impl fmt::Debug for Frame {

‎vm/src/vm.rs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -306,6 +306,10 @@ impl VirtualMachine {
306306
self.current_frame().last_value()
307307
}
308308

309+
fn nth_value(&self, i: usize) -> PyObjectRef {
310+
self.current_frame().nth_value(i)
311+
}
312+
309313
fn store_name(&mut self, name: &str) -> Option<PyResult> {
310314
let obj = self.pop_value();
311315
self.current_frame_mut().locals.set_item(name, obj);
@@ -890,6 +894,14 @@ impl VirtualMachine {
890894
self.push_value(obj);
891895
None
892896
}
897+
bytecode::Instruction::ListAppend { i } => {
898+
let list_obj = self.nth_value(*i);
899+
let item = self.pop_value();
900+
match self.call_method(&list_obj, "append", vec![item]) {
901+
Ok(_) => None,
902+
Err(err) => Some(Err(err)),
903+
}
904+
}
893905
bytecode::Instruction::BinaryOperation { ref op } => self.execute_binop(op),
894906
bytecode::Instruction::LoadAttr { ref name } => self.load_attr(name),
895907
bytecode::Instruction::StoreAttr { ref name } => self.store_attr(name),

0 commit comments

Comments
 (0)