Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit ebc723a

Browse files
committed
compile: implement comprehensions
1 parent 3c0ad78 commit ebc723a

File tree

3 files changed

+581
-5
lines changed

3 files changed

+581
-5
lines changed

compile/compile.go

+122-4
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,33 @@ func (c *compiler) compileAst(Ast ast.Ast, filename string, futureFlags int, don
188188
c.LoadConst(py.None)
189189
}
190190
c.Op(vm.RETURN_VALUE)
191+
case *ast.ListComp:
192+
// Elt Expr
193+
// Generators []Comprehension
194+
valueOnStack = true
195+
code.Name = "<listcomp>"
196+
c.OpArg(vm.BUILD_LIST, 0)
197+
c.comprehensionGenerator(node.Generators, 0, node.Elt, nil, Ast)
198+
case *ast.SetComp:
199+
// Elt Expr
200+
// Generators []Comprehension
201+
valueOnStack = true
202+
code.Name = "<setcomp>"
203+
c.OpArg(vm.BUILD_SET, 0)
204+
c.comprehensionGenerator(node.Generators, 0, node.Elt, nil, Ast)
205+
case *ast.DictComp:
206+
// Key Expr
207+
// Value Expr
208+
// Generators []Comprehension
209+
valueOnStack = true
210+
code.Name = "<dictcomp>"
211+
c.OpArg(vm.BUILD_MAP, 0)
212+
c.comprehensionGenerator(node.Generators, 0, node.Key, node.Value, Ast)
213+
case *ast.GeneratorExp:
214+
// Elt Expr
215+
// Generators []Comprehension
216+
code.Name = "<genexpr>"
217+
c.comprehensionGenerator(node.Generators, 0, node.Elt, nil, Ast)
191218

192219
default:
193220
panic(py.ExceptionNewf(py.SyntaxError, "Unknown ModuleBase: %v", Ast))
@@ -961,6 +988,97 @@ func (c *compiler) callHelper(n int, Args []ast.Expr, Keywords []*ast.Keyword, S
961988
c.OpArg(op, uint32(args+kwargs<<8))
962989
}
963990

991+
/* List and set comprehensions and generator expressions work by creating a
992+
nested function to perform the actual iteration. This means that the
993+
iteration variables don't leak into the current scope.
994+
The defined function is called immediately following its definition, with the
995+
result of that call being the result of the expression.
996+
The LC/SC version returns the populated container, while the GE version is
997+
flagged in symtable.c as a generator, so it returns the generator object
998+
when the function is called.
999+
This code *knows* that the loop cannot contain break, continue, or return,
1000+
so it cheats and skips the SETUP_LOOP/POP_BLOCK steps used in normal loops.
1001+
1002+
Possible cleanups:
1003+
- iterate over the generator sequence instead of using recursion
1004+
*/
1005+
func (c *compiler) comprehensionGenerator(generators []ast.Comprehension, gen_index int, elt ast.Expr, val ast.Expr, Ast ast.Ast) {
1006+
// generate code for the iterator, then each of the ifs,
1007+
// and then write to the element
1008+
start := new(Label)
1009+
skip := new(Label)
1010+
anchor := new(Label)
1011+
gen := generators[gen_index]
1012+
if gen_index == 0 {
1013+
/* Receive outermost iter as an implicit argument */
1014+
c.Code.Argcount = 1
1015+
c.OpArg(vm.LOAD_FAST, 0)
1016+
} else {
1017+
/* Sub-iter - calculate on the fly */
1018+
c.Expr(gen.Iter)
1019+
c.Op(vm.GET_ITER)
1020+
}
1021+
c.Label(start)
1022+
c.Jump(vm.FOR_ITER, anchor)
1023+
c.Expr(gen.Target)
1024+
1025+
/* XXX this needs to be cleaned up...a lot! */
1026+
for _, e := range gen.Ifs {
1027+
c.Expr(e)
1028+
c.Jump(vm.POP_JUMP_IF_FALSE, start)
1029+
}
1030+
1031+
gen_index++
1032+
if gen_index < len(generators) {
1033+
c.comprehensionGenerator(generators, gen_index, elt, val, Ast)
1034+
}
1035+
1036+
/* only append after the last for generator */
1037+
if gen_index >= len(generators) {
1038+
/* comprehension specific code */
1039+
switch Ast.(type) {
1040+
case *ast.GeneratorExp:
1041+
c.Expr(elt)
1042+
c.Op(vm.YIELD_VALUE)
1043+
c.Op(vm.POP_TOP)
1044+
case *ast.ListComp:
1045+
c.Expr(elt)
1046+
c.OpArg(vm.LIST_APPEND, uint32(gen_index+1))
1047+
case *ast.SetComp:
1048+
c.Expr(elt)
1049+
c.OpArg(vm.SET_ADD, uint32(gen_index+1))
1050+
case *ast.DictComp:
1051+
// With 'd[k] = v', v is evaluated before k, so we do the same.
1052+
c.Expr(val)
1053+
c.Expr(elt)
1054+
c.OpArg(vm.MAP_ADD, uint32(gen_index+1))
1055+
default:
1056+
panic(fmt.Sprintf("unknown comprehension %v", Ast))
1057+
}
1058+
c.Label(skip)
1059+
}
1060+
c.Jump(vm.JUMP_ABSOLUTE, start)
1061+
c.Label(anchor)
1062+
}
1063+
1064+
// Compile a comprehension
1065+
func (c *compiler) comprehension(expr ast.Expr, generators []ast.Comprehension) {
1066+
newSymTable := c.SymTable.FindChild(expr)
1067+
if newSymTable == nil {
1068+
panic("No symtable found for comprehension")
1069+
}
1070+
newC := newCompiler(c, compilerScopeComprehension)
1071+
code, err := newC.compileAst(expr, c.Code.Filename, 0, false, newSymTable)
1072+
if err != nil {
1073+
panic(err)
1074+
}
1075+
c.makeClosure(code, 0, newC, newC.Code.Name)
1076+
outermost_iter := generators[0].Iter
1077+
c.Expr(outermost_iter)
1078+
c.Op(vm.GET_ITER)
1079+
c.OpArg(vm.CALL_FUNCTION, 1)
1080+
}
1081+
9641082
// Compile expressions
9651083
func (c *compiler) Expr(expr ast.Expr) {
9661084
switch node := expr.(type) {
@@ -1078,20 +1196,20 @@ func (c *compiler) Expr(expr ast.Expr) {
10781196
case *ast.ListComp:
10791197
// Elt Expr
10801198
// Generators []Comprehension
1081-
panic("FIXME compile: ListComp not implemented")
1199+
c.comprehension(expr, node.Generators)
10821200
case *ast.SetComp:
10831201
// Elt Expr
10841202
// Generators []Comprehension
1085-
panic("FIXME compile: SetComp not implemented")
1203+
c.comprehension(expr, node.Generators)
10861204
case *ast.DictComp:
10871205
// Key Expr
10881206
// Value Expr
10891207
// Generators []Comprehension
1090-
panic("FIXME compile: DictComp not implemented")
1208+
c.comprehension(expr, node.Generators)
10911209
case *ast.GeneratorExp:
10921210
// Elt Expr
10931211
// Generators []Comprehension
1094-
panic("FIXME compile: GeneratorExp not implemented")
1212+
c.comprehension(expr, node.Generators)
10951213
case *ast.Yield:
10961214
// Value Expr
10971215
panic("FIXME compile: Yield not implemented")

0 commit comments

Comments
 (0)