rapydscript-ns 0.8.0

package/src/output/loops.pyj

@@ -0,0 +1,444 @@
+# vim:fileencoding=utf-8
+# License: BSD Copyright: 2016, Kovid Goyal <kovid at kovidgoyal.net>
+# globals: regenerate
+from __python__ import hash_literals
+
+from ast import AST_BaseCall, AST_SymbolRef, AST_Array, AST_Unary, AST_Number, has_calls, AST_Seq, AST_ListComprehension, AST_Starred, is_node_type
+from output.stream import OutputStream
+
+def unpack_tuple(elems, output, in_statement):
+    starred_idx = -1
+    for i, elem in enumerate(elems):
+        if is_node_type(elem, AST_Starred):
+            starred_idx = i
+            break
+    after_count = elems.length - starred_idx - 1 if starred_idx >= 0 else 0
+
+    for i, elem in enumerate(elems):
+        output.indent()
+        is_starred = is_node_type(elem, AST_Starred)
+        actual_elem = elem.expression if is_starred else elem
+        output.assign(actual_elem)
+        output.print("ρσ_unpack")
+        if is_starred:
+            output.print('.slice(')
+            output.print(i)
+            if after_count > 0:
+                output.print(', ρσ_unpack.length - ')
+                output.print(after_count)
+            output.print(')')
+        elif starred_idx >= 0 and i > starred_idx:
+            from_end = i - elems.length
+            output.print('[ρσ_unpack.length + ')
+            output.print(from_end)
+            output.print(']')
+        else:
+            output.with_square(def():
+                output.print(i)
+            )
+        if not in_statement or i < elems.length - 1:
+            output.semicolon()
+            output.newline()
+
+def print_do_loop(self, output):
+    output.print("do")
+    output.space()
+    self._do_print_body(output)
+    output.space()
+    output.print("while")
+    output.space()
+    output.with_parens(def(): self.condition.print(output);)
+    output.semicolon()
+
+def print_while_loop(self, output):
+    output.print("while")
+    output.space()
+    output.with_parens(def(): self.condition.print(output);)
+    output.space()
+    self._do_print_body(output)
+
+def is_simple_for_in(self):
+    # return true if this loop can be simplified into a basic for (i in j) loop
+    if is_node_type(self.object, AST_BaseCall)
+    and is_node_type(self.object.expression, AST_SymbolRef)
+    and self.object.expression.name is "dir" and self.object.args.length is 1:
+        return True
+    return False
+
+def is_simple_for(self):
+    # returns true if this loop can be simplified into a basic for(i=n;i<h;i++) loop
+    if (is_node_type(self.object, AST_BaseCall) and
+            is_node_type(self.object.expression, AST_SymbolRef) and
+            self.object.expression.name is "range" and
+            not (is_node_type(self.init, AST_Array))):
+        a = self.object.args
+        l = a.length
+        if l < 3 or (
+            is_node_type(a[2], AST_Number) or (
+                is_node_type(a[2], AST_Unary) and a[2].operator is '-' and is_node_type(a[2].expression, AST_Number)
+        )):
+            if (l is 1 and not has_calls(a[0])) or (l > 1 and not has_calls(a[1])):
+                return True
+    return False
+
+def print_for_loop_body(output):
+    self = this
+    output.with_block(def():
+        if not (self.simple_for_index or is_simple_for_in(self)):
+            # if we're using multiple iterators, unpack them
+            output.indent()
+            if output.options.js_version is 5:
+                itervar = "ρσ_Iter" + output.index_counter + "[ρσ_Index" + output.index_counter + "]"
+            else:
+                itervar = "ρσ_Index" + output.index_counter
+            if is_node_type(self.init, AST_Array):
+                flat = self.init.flatten()
+                output.assign("ρσ_unpack")
+                if flat.length > self.init.elements.length:
+                    output.print('ρσ_flatten(' + itervar + ')')
+                else:
+                    output.print(itervar)
+                output.end_statement()
+                unpack_tuple(flat, output)
+            else:
+                output.assign(self.init)
+                output.print(itervar)
+                output.end_statement()
+
+            output.index_counter += 1
+        if self.simple_for_index:
+            output.indent()
+            output.assign(self.init)
+            output.print(self.simple_for_index)
+            output.end_statement()
+
+        for stmt in self.body.body:
+            output.indent()
+            stmt.print(output)
+            output.newline()
+    )
+
+def init_es6_itervar(output, itervar):
+    output.indent()
+    output.spaced(itervar, '=', '((typeof', itervar + '[Symbol.iterator]', '===', '"function")', '?',
+                    '('+itervar, 'instanceof', 'Map', '?', itervar + '.keys()', ':', itervar+')', ':', 'Object.keys(' + itervar + '))')
+    output.end_statement()
+
+def print_for_in(self, output):
+    def write_object():
+        if self.object.constructor is AST_Seq:
+            (new AST_Array({'elements':self.object.to_array()})).print(output)
+        else:
+            self.object.print(output)
+
+    if is_simple_for(self):
+        # optimize range() into a simple for loop
+        increment = None
+        args = self.object.args
+        tmp_ = args.length
+        if tmp_ is 1:
+            start = 0
+            end = args[0]
+        elif tmp_ is 2:
+            start = args[0]
+            end = args[1]
+        elif tmp_ is 3:
+            start = args[0]
+            end = args[1]
+            increment = args[2]
+
+        self.simple_for_index = idx = 'ρσ_Index' + output.index_counter
+        output.index_counter += 1
+        output.print("for")
+        output.space()
+        output.with_parens(def():
+            output.spaced('var', idx, '='), output.space()
+            start.print(output) if start.print else output.print(start)
+            output.semicolon()
+            output.space()
+            output.print(idx)
+            output.space()
+            output.print(">") if is_node_type(increment, AST_Unary) else output.print("<")
+            output.space()
+            end.print(output)
+            output.semicolon()
+            output.space()
+            output.print(idx)
+            if increment and (not (is_node_type(increment, AST_Unary)) or increment.expression.value is not "1"):
+                if is_node_type(increment, AST_Unary):
+                    output.print("-=")
+                    increment.expression.print(output)
+                else:
+                    output.print("+=")
+                    increment.print(output)
+            else:
+                if is_node_type(increment, AST_Unary):
+                    output.print("--")
+                else:
+                    output.print("++")
+        )
+    elif is_simple_for_in(self):
+        # optimize dir() into a simple for in loop
+        output.print("for")
+        output.space()
+        output.with_parens(def():
+            self.init.print(output)
+            output.space()
+            output.print('in')
+            output.space()
+            self.object.args[0].print(output)
+        )
+    else:
+        # regular loop
+        if output.options.js_version is 5:
+            output.assign("var ρσ_Iter" + output.index_counter)
+            output.print("ρσ_Iterable")
+            output.with_parens(write_object)
+            output.semicolon()
+            output.newline()
+            output.indent()
+            output.print("for")
+            output.space()
+            output.with_parens(def():
+                output.print("var")
+                output.space()
+                output.assign("ρσ_Index" + output.index_counter)
+                output.print("0")
+                output.semicolon()
+                output.space()
+                output.print("ρσ_Index" + output.index_counter)
+                output.space()
+                output.print("<")
+                output.space()
+                output.print("ρσ_Iter" + output.index_counter + ".length")
+                output.semicolon()
+                output.space()
+                output.print("ρσ_Index" + output.index_counter + "++")
+            )
+        else:
+            itervar = "ρσ_Iter" + output.index_counter
+            output.assign("var " + itervar)
+            write_object()
+            output.end_statement()
+            init_es6_itervar(output, itervar)
+            output.indent()
+            output.spaced('for', '(var', 'ρσ_Index' + output.index_counter, 'of', itervar + ')')
+
+    output.space()
+    self._do_print_body(output)
+
+def print_list_comprehension(self, output):
+    tname = self.constructor.name.slice(4)
+    result_obj = {'ListComprehension':'[]', 'DictComprehension':('Object.create(null)' if self.is_jshash else '{}'), 'SetComprehension':'ρσ_set()'}[tname]
+    is_generator = tname is 'GeneratorComprehension'
+    es5 = output.options.js_version is 5
+    if tname is 'DictComprehension':
+        if self.is_pydict:
+            result_obj = 'ρσ_dict()'
+            add_to_result = def(output):
+                output.indent()
+                output.print('ρσ_Result.set')
+                output.with_parens(def():
+                    self.statement.print(output)
+                    output.space(), output.print(','), output.space()
+                    output.with_parens(def():
+                        if self.value_statement.constructor is AST_Seq:
+                            output.with_square(def():self.value_statement.print(output);)
+                        else:
+                            self.value_statement.print(output)
+                    )
+                )
+                output.end_statement()
+        else:
+            add_to_result = def(output):
+                output.indent()
+                output.print('ρσ_Result')
+                output.with_square(def():
+                    self.statement.print(output)
+                )
+                output.space(), output.print('='), output.space()
+                output.with_parens(def():
+                    if self.value_statement.constructor is AST_Seq:
+                        output.with_square(def():self.value_statement.print(output);)
+                    else:
+                        self.value_statement.print(output)
+                )
+                output.end_statement()
+    else:
+        push_func = "ρσ_Result." + ('push' if self.constructor is AST_ListComprehension else 'add')
+        if is_generator:
+            push_func = 'yield '
+        add_to_result = def(output):
+            output.indent()
+            output.print(push_func)
+            output.with_parens(def():
+                if self.statement.constructor is AST_Seq:
+                    output.with_square(def():self.statement.print(output);)
+                else:
+                    self.statement.print(output)
+            )
+            output.end_statement()
+
+    # Collect additional for-clauses (nested comprehension support)
+    clauses = self.clauses if self.clauses else []
+
+    # Helper: emit the for-loop for one clause level, recursing into inner clauses.
+    # depth 0 = outermost for-clause (uses ρσ_Iter / ρσ_Index),
+    # depth N = inner clause N (uses ρσ_IterN / ρσ_IndexN).
+    def emit_loop(body_out, depth):
+        if depth == 0:
+            iv = 'ρσ_Iter'
+            idx = 'ρσ_Index'
+            init_node = self.init
+            cond_node = self.condition
+        else:
+            iv = 'ρσ_Iter' + str(depth)
+            idx = 'ρσ_Index' + str(depth)
+            clause = clauses[depth - 1]
+            init_node = clause.init
+            cond_node = clause.condition
+            # Declare and initialise the inner iterable inside the outer loop body
+            body_out.indent()
+            body_out.assign("var " + iv)
+            if es5:
+                body_out.print("ρσ_Iterable")
+                body_out.with_parens(def(): clause.object.print(body_out);)
+            else:
+                clause.object.print(body_out)
+            body_out.end_statement()
+            if not es5:
+                init_es6_itervar(body_out, iv)
+        body_out.indent()
+        body_out.print("for")
+        body_out.space()
+        body_out.with_parens(def():
+            if es5:
+                body_out.print("var")
+                body_out.space()
+                body_out.assign(idx)
+                body_out.print("0")
+                body_out.semicolon()
+                body_out.space()
+                body_out.print(idx)
+                body_out.space()
+                body_out.print("<")
+                body_out.space()
+                body_out.print(iv + ".length")
+                body_out.semicolon()
+                body_out.space()
+                body_out.print(idx + "++")
+            else:
+                body_out.spaced('var', idx, 'of', iv)
+        )
+        body_out.space()
+        body_out.with_block(def():
+            body_out.indent()
+            itervar_expr = iv + '[' + idx + ']' if es5 else idx
+            if is_node_type(init_node, AST_Array):
+                flat = init_node.flatten()
+                body_out.assign("ρσ_unpack")
+                if flat.length > init_node.elements.length:
+                    body_out.print('ρσ_flatten(' + itervar_expr + ')')
+                else:
+                    body_out.print(itervar_expr)
+                body_out.end_statement()
+                unpack_tuple(flat, body_out)
+            else:
+                body_out.assign(init_node)
+                body_out.print(itervar_expr)
+                body_out.end_statement()
+            is_last = (depth == clauses.length)
+            if cond_node:
+                body_out.indent()
+                body_out.print("if")
+                body_out.space()
+                body_out.with_parens(def(): cond_node.print(body_out);)
+                body_out.space()
+                if is_last:
+                    body_out.with_block(def(): add_to_result(body_out);)
+                else:
+                    body_out.with_block(def(): emit_loop(body_out, depth + 1);)
+                body_out.newline()
+            else:
+                if is_last:
+                    add_to_result(body_out)
+                else:
+                    emit_loop(body_out, depth + 1)
+        )
+        body_out.newline()
+
+    output.with_parens(def():
+        output.print("function")
+        output.print("()")
+        output.space()
+        output.with_block(def():
+            body_out = output
+            if is_generator:
+                if es5:
+                    body_out = OutputStream({'beautify':True})
+                body_out.indent()
+                body_out.print('function* js_generator()'), body_out.space(), body_out.print('{')
+                body_out.newline()
+                previous_indentation = output.indentation()
+                output.set_indentation(output.next_indent())
+            body_out.indent()
+            body_out.assign("var ρσ_Iter")
+            if es5:
+                body_out.print("ρσ_Iterable")
+                body_out.with_parens(def():
+                    self.object.print(body_out)
+                )
+            else:
+                self.object.print(body_out)
+
+            if result_obj:
+                body_out.comma()
+                body_out.assign("ρσ_Result")
+                body_out.print(result_obj)
+            # Locally scope all loop variables (outer and inner clauses)
+            def decl_loop_vars(init_node):
+                if is_node_type(init_node, AST_Array):
+                    for i in init_node.elements:
+                        body_out.comma()
+                        i.print(body_out)
+                else:
+                    body_out.comma()
+                    init_node.print(body_out)
+            decl_loop_vars(self.init)
+            for clause in clauses:
+                decl_loop_vars(clause.init)
+            body_out.end_statement()
+
+            if not es5:
+                init_es6_itervar(body_out, 'ρσ_Iter')
+            emit_loop(body_out, 0)
+            if self.constructor is AST_ListComprehension:
+                body_out.indent()
+                body_out.spaced('ρσ_Result', '=', 'ρσ_list_constructor(ρσ_Result)')
+                body_out.end_statement()
+            if not is_generator:
+                body_out.indent()
+                body_out.print("return ρσ_Result")
+                body_out.end_statement()
+            if is_generator:
+                output.set_indentation(previous_indentation)
+                body_out.newline(), body_out.indent(), body_out.print('}')  # end js_generator
+                if es5:
+                    transpiled = regenerate(body_out.get(), output.options.beautify).replace(/regeneratorRuntime.(wrap|mark)/g, 'ρσ_regenerator.regeneratorRuntime.$1')
+                    if output.options.beautify:
+                        ci = output.make_indent(0)
+                        transpiled = [ci + x for x in transpiled.split('\n')].join('\n')
+                    output.print(transpiled)
+                output.newline(), output.indent()
+                output.spaced('var', 'result', '=', 'js_generator.call(this)')
+                output.end_statement()
+                # Python's generator objects use a separate method to send data to the generator
+                output.indent()
+                output.spaced('result.send', '=', 'result.next')
+                output.end_statement()
+                output.indent()
+                output.spaced('return', 'result')
+                output.end_statement()
+        )
+    )
+    output.print("()")