<commit_before># Copyright (c) 2012-2014 The Mirah project authors. All Rights Reserved.
# All contributing project authors may be found in the NOTICE file.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

package org.mirah.typer

import mirah.lang.ast.*
import java.util.logging.Level
import org.mirah.util.Logger
import java.util.Collections
import java.util.Collection
import java.util.LinkedHashMap
import java.util.IdentityHashMap
import java.util.HashSet
import java.util.LinkedHashSet
import java.util.List
import java.util.Stack
import java.util.Map
import java.util.Map.Entry
import java.util.ArrayList
import java.io.File

import org.mirah.jvm.compiler.ProxyCleanup
import org.mirah.jvm.mirrors.MirrorScope
import org.mirah.jvm.mirrors.BaseType
import org.mirah.jvm.mirrors.MirrorTypeSystem
import org.mirah.jvm.mirrors.MirrorFuture
import org.mirah.jvm.mirrors.MethodScope
import org.mirah.jvm.mirrors.ClosureScope
import org.mirah.macros.MacroBuilder
import org.mirah.typer.simple.TypePrinter2
import org.mirah.typer.CallFuture
import org.mirah.typer.BaseTypeFuture
import org.mirah.util.AstFormatter
import org.mirah.util.AstChecker


interface ClosureBuilderer
  def insert_closure(block: Block, parent_type: ResolvedType): TypeFuture; end
  def add_todo(block: Block, parent_type: ResolvedType): void; end
  def finish: void; end
end


# This class transforms a Block into an anonymous class once the Typer has figured out
# the interface to implement (or the abstract superclass).
#
# Note: This is ugly. It depends on the internals of the JVM scope and jvm_bytecode classes,
# and the BindingReference node is a hack. This should really all be cleaned up.


# better idea:
#   add_todo doesn't add blocks + types.
#      it adds Script parents to set of them
#      finish iters over scripts to find block -> type map
#      then does them
class BetterClosureBuilder
  implements ClosureBuilderer

  attr_reader typer: Typer

  def self.initialize: void
    @@log = Logger.getLogger(BetterClosureBuilder.class.getName)
  end

  def initialize(typer: Typer, macros: MacroBuilder)
    @typer = typer
    @types = typer.type_system
    @scoper = typer.scoper
    
    @todo_closures = LinkedHashMap.new
    @scripts = LinkedHashSet.new
    @macros = macros
  end


# NLR
# in finish
# iter blocks
#  check for returns
#  next if !returns
#  gen up exception type
#  replace returns with raise of exception type
#  re-infer
#  enclosing = find Script or MethodDefinition
#  save (exceptiontype, enclosing) tuple
# iter enclosings
#  wrap w/ rescue of exceptiontype
#  re-infer


  # TODO better positions!!
  # the positions should be synthetic, or should be carefully pointed at the closest thing
  # maybe source locations could encode some info...
  class BindingAdjuster < NodeScanner
    def self.initialize
      @@log = Logger.getLogger(BindingAdjuster.class.getName)
    end

    def initialize(
       builder: BetterClosureBuilder,
       bindingName: String,
       parent_scope: MirrorScope,
       blockToBindings: Map,
       bindingLocalNamesToTypes: Map)

      @builder = builder

      @blockToBindings = blockToBindings
      @bindingLocalNamesToTypes = bindingLocalNamesToTypes

      @blocks = Stack.new
      @parent_scope = parent_scope
      @bindingName = bindingName
    end

    def adjust(node: Node, block: Block): boolean
      @captured = @parent_scope.capturedLocals
      @@log.fine "adjusting #{node}\n#{@builder.typer.sourceContent node}\nfor block\n#{@builder.typer.sourceContent block}"
      @@log.fine "captures for #{@bindingName}: #{@captured} parent scope: #{@parent_scope}"

      if @captured.isEmpty
        @@log.fine "no need for binding adjustment here. Nothing captured"
        return false
      end
      if @parent_scope.declared_binding_type
        @@log.fine "no need for binding adjustment here. already bound to #{@parent_scope.declared_binding_type}"
        return false
      end

      # construct binding
      name = @builder.temp_name_from_outer_scope(node, "ZBinding")
      
      @@log.fine("building binding #{name} with captures #{@captured}")
      binding_klass = @builder.build_class(
        node.position, nil, name)

      entries = @captured.map do |cap: String|
        type = @parent_scope.getLocalType(cap, node.position).resolve
        HashEntry.new(SimpleString.new(cap), Constant.new(SimpleString.new(type.name)))
      end

      attr_def = FunctionalCall.new(
        SimpleString.new("attr_accessor"),
        [Hash.new(node.position, entries)],
        nil)
      binding_klass.body.insert(0, attr_def)
            
      binding_new_call = Call.new(node.position, Constant.new(SimpleString.new(name)), SimpleString.new("new"), [], nil)
      @builder.typer.workaroundASTBug binding_new_call

      assign_binding_dot_new = LocalAssignment.new(
          node.position,
          SimpleString.new(@bindingName),
          binding_new_call)
      
      @@log.fine "inserted binding assign / binding class "
      @builder.insert_into_body NodeList(node), assign_binding_dot_new
      @builder.insert_into_body NodeList(node), binding_klass
      
      binding_type = @builder.infer(binding_klass).resolve

      raise "parent_scope had declared_binding_type already #{@parent_scope}" if @parent_scope.declared_binding_type
      @parent_scope.declared_binding_type = binding_type
      @bindingLocalNamesToTypes[@bindingName] = binding_type
      @builder.parent_scope_to_binding_name[@parent_scope] = @bindingName
      
      @binding_type = @builder.infer(binding_klass)
      @binding_klass_node = binding_klass
      @builder.infer assign_binding_dot_new
      @@log.fine "binding assignment inference done"

      @@log.fine "replacing references to captures"

      mdef = node.findAncestor{ |n| n.kind_of? MethodDefinition }
      block_parent = node.findAncestor { |n| n.kind_of? Block }

      arguments = if mdef
        MethodDefinition(mdef).arguments
      elsif block_parent
        Block(block_parent).arguments
      end

      node.accept self, 1
      @@log.fine "finished phase one of capture replacement"

      arg_names = []
      if arguments
        arguments.required.each {|a: FormalArgument| arg_names.add a.name.identifier } if arguments.required
        arguments.optional.each {|a: FormalArgument| arg_names.add a.name.identifier } if arguments.optional
        arg_names.add arguments.rest.name.identifier if arguments.rest
        arguments.required2.each {|a: FormalArgument| arg_names.add a.name.identifier } if arguments.required2
        arg_names.add arguments.block.name.identifier if arguments.block
      else
        if mdef || block_parent
          @@log.fine "parent had no arguments: parent #{mdef} #{block_parent}"
        else
          @@log.fine "had no parent"
        end
      end
      @@log.fine "adding assignments from args to captures"
      arg_names.each do |arg|
        if @captured.contains arg
          addition = Call.new(
            blockAccessNode(node.position),
            SimpleString.new("#{arg}_set"),
            [LocalAccess.new(node.position, SimpleString.new(arg))],
            nil)
          @builder.typer.workaroundASTBug addition
           # insert after binding class & binding instantiation
           # if we were less lazy, we'd find the binding construction, and insert after
          NodeList(node).insert 2, addition

          @builder.infer addition
        end
      end

      @@log.fine "done replacing references"
      true
    end

    def visitClassDefinition(node, blah)
      nil # do not descent into classes which just happen to be defined in the scope which gets a closure
    end
    
    def enterClosureDefinition(node, blah)
      # skip the definition of the binding we're in the process of inserting
      @binding_klass_node != node
    end


    def exitClosureDefinition(node, blah)
      # might be a lambda created by the lambda macro
      # if contained captures?
      #  ; containing captures here implies a closure not built via this, because this builds them inside out
      #    adjust initializer to add binding
      #    add method for referring to binding by name
      #    find new call, and adjust it to include the binding just added to the initializer
    end

    def enterBlock(node, blah)
      @blocks.push node
      @blockToBindings[@builder.blockCloneMapNewOld[@blocks.peek]] ||= HashSet.new

      true
    end

    def exitBlock(node, blah)
      @blocks.pop
    end

    def maybeNoteBlockBinding
      @blocks.each do |block: Block|
        Collection(@blockToBindings[@builder.blockCloneMapNewOld[block]]).add @bindingName
      end
    end

    def blockAccessNode(position: Position)
      name_node = SimpleString.new(@bindingName)
      if @blocks.isEmpty
        LocalAccess.new(position, name_node)
      else
        FieldAccess.new(position, name_node)
      end
    end

    def exitLocalAssignment(local, blah)
      local_name = local.name.identifier
      return nil unless @captured.contains local_name

      @@log.finest "enterLocalAssignment: replacing #{local.name.identifier} with #{@bindingName}.#{local.name.identifier}="
      @@log.finest "  Type: #{@builder.typer.getInferredType(local)}"

      maybeNoteBlockBinding

      new_value = Node(local.value)
      new_value.setParent(nil)
      replacement = Call.new(
        blockAccessNode(local.position),
        SimpleString.new("#{local.name.identifier}_set"),
        [new_value],
        nil)

      @builder.typer.workaroundASTBug replacement

      replaceSelf(local, replacement)
      local.value.setParent replacement

      @builder.typer.learnType replacement.target, @binding_type
      @builder.typer.infer new_value
      @builder.typer.infer replacement
    end

    def exitLocalAccess(local, blah)
      local_name = local.name.identifier
      return nil unless @captured.contains local_name

      @@log.finest "enterLocalAccess: replacing #{local.name.identifier} with #{@bindingName}.#{local.name.identifier}="
      @@log.finest "  Type: #{@builder.typer.getInferredType(local)}"
      
      maybeNoteBlockBinding

      replacement = Call.new(
        blockAccessNode(local.position),
        SimpleString.new(local.position,local.name.identifier),
        [],
        nil)

      @builder.typer.workaroundASTBug replacement

      replaceSelf(local, replacement)

      @builder.typer.learnType replacement.target, @binding_type

      @@log.fine "does replacement have parent?: #{replacement.parent}"
      @@log.fine "does replacement parent 0th child: #{NodeList(replacement.parent).get(0)}" if replacement.parent.kind_of? NodeList
      @@log.fine "does replacement have parent?: #{replacement.parent.parent}" if replacement.parent
      @builder.typer.infer replacement
    end

    def replaceSelf me: Node, replacement: Node
      me.parent.replaceChild(me, replacement)
    end
  end

  class BlockFinder < NodeScanner
    def initialize(typer: Typer, todo_closures: Map)
      @typer = typer
      @todo_closures = todo_closures
    end
    def find(node: Script): Map
      collection = LinkedHashMap.new
      node.accept(self, collection)
      collection
    end

    def enterMacroDefinition(node, notes)
      false
    end

    def exitBlock(node, notes)
      type = if @todo_closures[node]
        @todo_closures[node]
      else
        #t = @typer.getInferredType(node).resolve
        #if t.kind_of? MethodType
          parent_type_future = @typer.getInferredType(node.parent)
          unless parent_type_future
            puts "#{CallSite(node.parent).name} call with block has no type at #{node.parent.position}"
            puts "  block type: #{@typer.getInferredType(node)}"
          end
          fs = CallFuture(parent_type_future).futures
          TypeFuture(fs.get(fs.size-1)).resolve
        #else
        #end
      end
      Map(notes).put node, type
    end
  end

  attr_accessor blockCloneMapOldNew:IdentityHashMap
  attr_accessor blockCloneMapNewOld:IdentityHashMap
  attr_accessor parent_scope_to_binding_name:Map

  def finish
    closures = []
    scripts = ArrayList.new(@scripts)
    Collections.reverse(scripts)
    scripts.each do |s: Script|
      closures.addAll BlockFinder.new(@typer, @todo_closures).find(s).entrySet
    end

    closures_to_skip = []

    blockToBindings = LinkedHashMap.new # the list of bindings a block closes over
    bindingLocalNamesToTypes = LinkedHashMap.new

    bindingForBlocks = LinkedHashMap.new # the specific binding for a given block
    Collections.reverse(closures) # from outside to inside
    
    self.blockCloneMapOldNew = IdentityHashMap.new
    self.blockCloneMapNewOld = IdentityHashMap.new
    
    selff = self
    
    closures.each do |entry: Entry|
      block = Block(entry.getKey)
      block.whenCloned do |interim,new|
        old = selff.blockCloneMapNewOld.get(interim)
        selff.blockCloneMapNewOld.put(new,old)
#       blockCloneMapOldNew.remove(old)
        selff.blockCloneMapOldNew.put(old,new)
      end
      blockCloneMapOldNew.put(block,block)
      blockCloneMapNewOld.put(block,block)
    end

    self.parent_scope_to_binding_name = {}

    i = 0
    closures.each do |entry: Entry|
      @@log.fine "adjust bindings for block #{entry.getKey} #{entry.getValue} #{i}"
      i += 1
      uncloned_block = Block(entry.getKey)
      block = Block(blockCloneMapOldNew.get(uncloned_block))
      @@log.fine "#{typer.sourceContent block}"
      enclosing_node = find_enclosing_node(block)
      if enclosing_node.nil?
        # this likely means a macro exists and made things confusing 
        # by copying the tree
        @@log.fine "enclosing node was nil, removing  #{entry.getKey} #{entry.getValue} #{i}"
        closures_to_skip.add entry
        next
      end
      @@log.fine "enclosing node #{enclosing_node}"
      @@log.fine "#{typer.sourceContent enclosing_node}"

      ProxyCleanup.new.scan enclosing_node

      enclosing_b = get_body(enclosing_node)      
      if enclosing_b.nil?
        closures_to_skip.add entry
        next
      end
      bindingName = "b#{i}"
      bindingForBlocks.put uncloned_block, bindingName
      parent_scope = MirrorScope(get_scope(block)) 
      adjuster = BindingAdjuster.new(
        self,
        bindingName,
        parent_scope,
        blockToBindings,
        bindingLocalNamesToTypes)

      binding_generated = adjuster.adjust enclosing_b, block
      @@log.log(Level.FINE, "After adjusting: #{AstFormatter.new(Node(scripts.get(0)))}")
#    end
#
#    # ignore closures with no parents, they aren't in the final AST, maybe
#    closures.removeAll closures_to_skip
#
#    i = 0
#    closures.each do |entry: Entry|
#      @@log.fine "insert_closure #{entry.getKey} #{entry.getValue} #{i}"
#      i += 1

      if AstChecker.enabled
        scripts.each do |s: Script|
          s.accept(AstChecker.new,nil)
        end
      end
  
      block = Block(blockCloneMapOldNew.get(entry.getKey))
      parent_type = ResolvedType(entry.getValue)

      unless get_body(find_enclosing_node(block))
        @@log.fine "  enclosing node was nil, removing  #{entry.getKey} #{entry.getValue} #{i}"
        next
      end

      closure_name = temp_name_from_outer_scope(block, "Closure")
      closure_klass = build_class(block.position, parent_type, closure_name)

      # build closure class
      binding_list = Collection(blockToBindings.get(uncloned_block)) || Collections.emptyList
      binding_args = binding_list.map do |name: String|
        RequiredArgument.new(SimpleString.new(name), SimpleString.new(ResolvedType(bindingLocalNamesToTypes[name]).name))
      end

      args = Arguments.new(closure_klass.position,
                           binding_args,
                           Collections.emptyList,
                           nil,
                           Collections.emptyList,
                           nil)
      binding_assigns = binding_list.map do |name: String|
        FieldAssign.new(SimpleString.new(name), LocalAccess.new(SimpleString.new(name)), nil)
      end
      constructor = ConstructorDefinition.new(
        SimpleString.new('initialize'), args,
        SimpleString.new('void'), binding_assigns, nil)
      closure_klass.body.add(constructor)

      enclosing_b  = find_enclosing_body block
      if binding_assigns.isEmpty
        insert_into_body enclosing_b, closure_klass
      else 
        # insert after binding class, for happier typing
        # also causes weird issues in the compiler if things are out of order
        enclosing_b.insert(1, closure_klass)
      end

      block_scope = get_scope block
      if contains_methods(block)
        copy_methods(closure_klass, block, block_scope)
      else
        build_and_inject_methods(closure_klass, block, parent_type, block_scope)
      end

      closure_type = infer(closure_klass) # FIXME: this re-infers also the body of the method (which is the ex-body of the block), which is probably duplicate work.

      has_block_parent = block.findAncestor { |node| node.parent.kind_of?(Block) || node.parent.kind_of?(ClosureDefinition) } # block, or converted block

      binding_locals = binding_list.map do |name: String|
        if has_block_parent && !name.equals(parent_scope_to_binding_name[parent_scope])
          FieldAccess.new(SimpleString.new(name))
        else
          LocalAccess.new(SimpleString.new(name))
        end
      end
      target = makeTypeName(block.position, closure_type.resolve)
      new_node = Call.new(
        block.position, target,
        SimpleString.new("new"), 
        binding_locals, nil)
      @typer.workaroundASTBug new_node

      

      if block.parent.kind_of?(CallSite)
        parent = CallSite(block.parent)
        replace_block_with_closure_in_call parent, block, new_node
      else
        replace_synthetic_lambda_definiton_with_closure(SyntheticLambdaDefinition(block.parent),new_node)
      end

      @@log.fine "inferring new_node #{new_node}"
      infer new_node
      @@log.fine "inferring enclosing_b #{enclosing_b}"
      infer enclosing_b

      @@log.fine "done with #{enclosing_b}"
      @@log.log(Level.FINE, "Inferred AST:\n{0}", AstFormatter.new(enclosing_b))

      buf = java::io::ByteArrayOutputStream.new
      ps = java::io::PrintStream.new(buf)
      printer = TypePrinter2.new(@typer, ps)
      printer.scan(enclosing_b, nil)
      ps.close()
      @@log.fine("Inferred types for expr:\n#{String.new(buf.toByteArray)}")
    end
  end

  def add_todo(block: Block, parent_type: ResolvedType)
    return if parent_type.isError || block.parent.nil?
    
    rtype = BaseTypeFuture.new(block.position)
    rtype.resolved parent_type
  
    new_scope = @typer.addNestedScope block
    new_scope.selfType = rtype
    if contains_methods block
      @typer.infer block.body
    else
      @typer.inferClosureBlock block, method_for(parent_type)
    end

    script = block.findAncestor{|n| n.kind_of? Script}

    @todo_closures[block] = parent_type
    @scripts.add script
  end

  def insert_closure(block: Block, parent_type: ResolvedType)
    raise "BetterClosureBuilder doesn't support insert_closure"
  end

  def prepare_non_local_return_closure(block: Block, parent_type: ResolvedType): Node
    # generates closure classes, AND an exception type
    # and replaces the closure call with something like this:
    #
    # class MyNonLocalReturn < Throwable
    #   def initialize(return_value:`method return type`); @return_value = return_value; end
    #   def return_value; @return_value; end
    # end
    # begin
    #   call { raise MyNonLocalReturn, `value` }
    # rescue MyNonLocalReturn => e
    #   return e.return_value
    # end
    enclosing_node = find_enclosing_node block
    return_type = if enclosing_node.kind_of? MethodDefinition
                    methodType = infer(enclosing_node)
                    MethodFuture(methodType).returnType
                  elsif enclosing_node.kind_of? Script
                    future = AssignableTypeFuture.new block.position
                    future.assign(infer(enclosing_node), block.position)
                    future
                  end
    nlr_klass = define_nlr_exception block
    block = convert_returns_to_raises block, nlr_klass, return_type
    new_node = nlr_prepare block, parent_type, nlr_klass
    resolved = return_type.resolve

    raise "Unable to determine method return type before generating closure including non local return" unless resolved

    enclosing_body = get_body(enclosing_node)
    node_in_body = block.findAncestor { |node| node.parent.kind_of? NodeList }
    new_call = wrap_with_rescue block, nlr_klass, node_in_body, resolved
    node_in_body.parent.replaceChild node_in_body, new_call
    
    finish_nlr_exception block, nlr_klass, resolved
    insert_into_body enclosing_body, nlr_klass
    infer(nlr_klass)
    new_node
  end

  # creates closure class, inserts it
  # returns Call Node that is the instantiation of the closure
  def prepare(block: Block, parent_type: ResolvedType): Call
    Call(prepare_regular_closure(block, parent_type))
  end

  def prepare_regular_closure(block: Block, parent_type: ResolvedType): Node
    parent_scope = get_scope block
    klass = build_closure_class block, parent_type, parent_scope
    if contains_methods(block)
      copy_methods(klass, block, parent_scope)
    else
      build_and_inject_methods(klass, block, parent_type, parent_scope)
    end
    new_closure_call_node(block, klass)
  end

  def replace_block_with_closure_in_call(parent: CallSite, block: Block, new_node: Node): void
    if block == parent.block
      parent.block = nil
      parent.parameters.add(new_node)
    else
      new_node.setParent(nil)
      parent.replaceChild(block, new_node)
    end
  end
  
  def replace_synthetic_lambda_definiton_with_closure(parent: SyntheticLambdaDefinition, new_node: Node): void
    parentparent = parent.parent
    new_node.setParent(nil)
    if parentparent.kind_of?(CallSite) # then the SyntheticLambdaDefinition is not a child of the CallSite itself, but (most likely?) a child of its arguments. FIXME: It is weird that the parent of a child of X is not X. 
      CallSite(parentparent).parameters.replaceChild(parent,new_node)
    else
      parentparent.replaceChild(parent,new_node)
    end
  end

  def find_enclosing_body block: Block
    enclosing_node = find_enclosing_node block
    get_body enclosing_node
  end

  def find_enclosing_method_body block: Block
    enclosing_node = find_enclosing_method block
    get_body enclosing_node
  end

  def get_body node: Node
    # TODO create an interface for nodes with bodies
    if node.kind_of?(MethodDefinition)
      MethodDefinition(node).body
    elsif node.kind_of?(Script)
      Script(node).body
    elsif node.kind_of?(Block)
      Block(node).body
    else
      raise "Unknown type for finding a body #{node.getClass}"
    end
  end

  def find_enclosing_node block: Node
    if block.parent<commit_msg>Fix case where SyntheticLambda is a call target.<commit_after># Copyright (c) 2012-2014 The Mirah project authors. All Rights Reserved.
# All contributing project authors may be found in the NOTICE file.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

package org.mirah.typer

import mirah.lang.ast.*
import java.util.logging.Level
import org.mirah.util.Logger
import java.util.Collections
import java.util.Collection
import java.util.LinkedHashMap
import java.util.IdentityHashMap
import java.util.HashSet
import java.util.LinkedHashSet
import java.util.List
import java.util.Stack
import java.util.Map
import java.util.Map.Entry
import java.util.ArrayList
import java.io.File

import org.mirah.jvm.compiler.ProxyCleanup
import org.mirah.jvm.mirrors.MirrorScope
import org.mirah.jvm.mirrors.BaseType
import org.mirah.jvm.mirrors.MirrorTypeSystem
import org.mirah.jvm.mirrors.MirrorFuture
import org.mirah.jvm.mirrors.MethodScope
import org.mirah.jvm.mirrors.ClosureScope
import org.mirah.macros.MacroBuilder
import org.mirah.typer.simple.TypePrinter2
import org.mirah.typer.CallFuture
import org.mirah.typer.BaseTypeFuture
import org.mirah.util.AstFormatter
import org.mirah.util.AstChecker


interface ClosureBuilderer
  def insert_closure(block: Block, parent_type: ResolvedType): TypeFuture; end
  def add_todo(block: Block, parent_type: ResolvedType): void; end
  def finish: void; end
end


# This class transforms a Block into an anonymous class once the Typer has figured out
# the interface to implement (or the abstract superclass).
#
# Note: This is ugly. It depends on the internals of the JVM scope and jvm_bytecode classes,
# and the BindingReference node is a hack. This should really all be cleaned up.


# better idea:
#   add_todo doesn't add blocks + types.
#      it adds Script parents to set of them
#      finish iters over scripts to find block -> type map
#      then does them
class BetterClosureBuilder
  implements ClosureBuilderer

  attr_reader typer: Typer

  def self.initialize: void
    @@log = Logger.getLogger(BetterClosureBuilder.class.getName)
  end

  def initialize(typer: Typer, macros: MacroBuilder)
    @typer = typer
    @types = typer.type_system
    @scoper = typer.scoper
    
    @todo_closures = LinkedHashMap.new
    @scripts = LinkedHashSet.new
    @macros = macros
  end


# NLR
# in finish
# iter blocks
#  check for returns
#  next if !returns
#  gen up exception type
#  replace returns with raise of exception type
#  re-infer
#  enclosing = find Script or MethodDefinition
#  save (exceptiontype, enclosing) tuple
# iter enclosings
#  wrap w/ rescue of exceptiontype
#  re-infer


  # TODO better positions!!
  # the positions should be synthetic, or should be carefully pointed at the closest thing
  # maybe source locations could encode some info...
  class BindingAdjuster < NodeScanner
    def self.initialize
      @@log = Logger.getLogger(BindingAdjuster.class.getName)
    end

    def initialize(
       builder: BetterClosureBuilder,
       bindingName: String,
       parent_scope: MirrorScope,
       blockToBindings: Map,
       bindingLocalNamesToTypes: Map)

      @builder = builder

      @blockToBindings = blockToBindings
      @bindingLocalNamesToTypes = bindingLocalNamesToTypes

      @blocks = Stack.new
      @parent_scope = parent_scope
      @bindingName = bindingName
    end

    def adjust(node: Node, block: Block): boolean
      @captured = @parent_scope.capturedLocals
      @@log.fine "adjusting #{node}\n#{@builder.typer.sourceContent node}\nfor block\n#{@builder.typer.sourceContent block}"
      @@log.fine "captures for #{@bindingName}: #{@captured} parent scope: #{@parent_scope}"

      if @captured.isEmpty
        @@log.fine "no need for binding adjustment here. Nothing captured"
        return false
      end
      if @parent_scope.declared_binding_type
        @@log.fine "no need for binding adjustment here. already bound to #{@parent_scope.declared_binding_type}"
        return false
      end

      # construct binding
      name = @builder.temp_name_from_outer_scope(node, "ZBinding")
      
      @@log.fine("building binding #{name} with captures #{@captured}")
      binding_klass = @builder.build_class(
        node.position, nil, name)

      entries = @captured.map do |cap: String|
        type = @parent_scope.getLocalType(cap, node.position).resolve
        HashEntry.new(SimpleString.new(cap), Constant.new(SimpleString.new(type.name)))
      end

      attr_def = FunctionalCall.new(
        SimpleString.new("attr_accessor"),
        [Hash.new(node.position, entries)],
        nil)
      binding_klass.body.insert(0, attr_def)
            
      binding_new_call = Call.new(node.position, Constant.new(SimpleString.new(name)), SimpleString.new("new"), [], nil)
      @builder.typer.workaroundASTBug binding_new_call

      assign_binding_dot_new = LocalAssignment.new(
          node.position,
          SimpleString.new(@bindingName),
          binding_new_call)
      
      @@log.fine "inserted binding assign / binding class "
      @builder.insert_into_body NodeList(node), assign_binding_dot_new
      @builder.insert_into_body NodeList(node), binding_klass
      
      binding_type = @builder.infer(binding_klass).resolve

      raise "parent_scope had declared_binding_type already #{@parent_scope}" if @parent_scope.declared_binding_type
      @parent_scope.declared_binding_type = binding_type
      @bindingLocalNamesToTypes[@bindingName] = binding_type
      @builder.parent_scope_to_binding_name[@parent_scope] = @bindingName
      
      @binding_type = @builder.infer(binding_klass)
      @binding_klass_node = binding_klass
      @builder.infer assign_binding_dot_new
      @@log.fine "binding assignment inference done"

      @@log.fine "replacing references to captures"

      mdef = node.findAncestor{ |n| n.kind_of? MethodDefinition }
      block_parent = node.findAncestor { |n| n.kind_of? Block }

      arguments = if mdef
        MethodDefinition(mdef).arguments
      elsif block_parent
        Block(block_parent).arguments
      end

      node.accept self, 1
      @@log.fine "finished phase one of capture replacement"

      arg_names = []
      if arguments
        arguments.required.each {|a: FormalArgument| arg_names.add a.name.identifier } if arguments.required
        arguments.optional.each {|a: FormalArgument| arg_names.add a.name.identifier } if arguments.optional
        arg_names.add arguments.rest.name.identifier if arguments.rest
        arguments.required2.each {|a: FormalArgument| arg_names.add a.name.identifier } if arguments.required2
        arg_names.add arguments.block.name.identifier if arguments.block
      else
        if mdef || block_parent
          @@log.fine "parent had no arguments: parent #{mdef} #{block_parent}"
        else
          @@log.fine "had no parent"
        end
      end
      @@log.fine "adding assignments from args to captures"
      arg_names.each do |arg|
        if @captured.contains arg
          addition = Call.new(
            blockAccessNode(node.position),
            SimpleString.new("#{arg}_set"),
            [LocalAccess.new(node.position, SimpleString.new(arg))],
            nil)
          @builder.typer.workaroundASTBug addition
           # insert after binding class & binding instantiation
           # if we were less lazy, we'd find the binding construction, and insert after
          NodeList(node).insert 2, addition

          @builder.infer addition
        end
      end

      @@log.fine "done replacing references"
      true
    end

    def visitClassDefinition(node, blah)
      nil # do not descent into classes which just happen to be defined in the scope which gets a closure
    end
    
    def enterClosureDefinition(node, blah)
      # skip the definition of the binding we're in the process of inserting
      @binding_klass_node != node
    end


    def exitClosureDefinition(node, blah)
      # might be a lambda created by the lambda macro
      # if contained captures?
      #  ; containing captures here implies a closure not built via this, because this builds them inside out
      #    adjust initializer to add binding
      #    add method for referring to binding by name
      #    find new call, and adjust it to include the binding just added to the initializer
    end

    def enterBlock(node, blah)
      @blocks.push node
      @blockToBindings[@builder.blockCloneMapNewOld[@blocks.peek]] ||= HashSet.new

      true
    end

    def exitBlock(node, blah)
      @blocks.pop
    end

    def maybeNoteBlockBinding
      @blocks.each do |block: Block|
        Collection(@blockToBindings[@builder.blockCloneMapNewOld[block]]).add @bindingName
      end
    end

    def blockAccessNode(position: Position)
      name_node = SimpleString.new(@bindingName)
      if @blocks.isEmpty
        LocalAccess.new(position, name_node)
      else
        FieldAccess.new(position, name_node)
      end
    end

    def exitLocalAssignment(local, blah)
      local_name = local.name.identifier
      return nil unless @captured.contains local_name

      @@log.finest "enterLocalAssignment: replacing #{local.name.identifier} with #{@bindingName}.#{local.name.identifier}="
      @@log.finest "  Type: #{@builder.typer.getInferredType(local)}"

      maybeNoteBlockBinding

      new_value = Node(local.value)
      new_value.setParent(nil)
      replacement = Call.new(
        blockAccessNode(local.position),
        SimpleString.new("#{local.name.identifier}_set"),
        [new_value],
        nil)

      @builder.typer.workaroundASTBug replacement

      replaceSelf(local, replacement)
      local.value.setParent replacement

      @builder.typer.learnType replacement.target, @binding_type
      @builder.typer.infer new_value
      @builder.typer.infer replacement
    end

    def exitLocalAccess(local, blah)
      local_name = local.name.identifier
      return nil unless @captured.contains local_name

      @@log.finest "enterLocalAccess: replacing #{local.name.identifier} with #{@bindingName}.#{local.name.identifier}="
      @@log.finest "  Type: #{@builder.typer.getInferredType(local)}"
      
      maybeNoteBlockBinding

      replacement = Call.new(
        blockAccessNode(local.position),
        SimpleString.new(local.position,local.name.identifier),
        [],
        nil)

      @builder.typer.workaroundASTBug replacement

      replaceSelf(local, replacement)

      @builder.typer.learnType replacement.target, @binding_type

      @@log.fine "does replacement have parent?: #{replacement.parent}"
      @@log.fine "does replacement parent 0th child: #{NodeList(replacement.parent).get(0)}" if replacement.parent.kind_of? NodeList
      @@log.fine "does replacement have parent?: #{replacement.parent.parent}" if replacement.parent
      @builder.typer.infer replacement
    end

    def replaceSelf me: Node, replacement: Node
      me.parent.replaceChild(me, replacement)
    end
  end

  class BlockFinder < NodeScanner
    def initialize(typer: Typer, todo_closures: Map)
      @typer = typer
      @todo_closures = todo_closures
    end
    def find(node: Script): Map
      collection = LinkedHashMap.new
      node.accept(self, collection)
      collection
    end

    def enterMacroDefinition(node, notes)
      false
    end

    def exitBlock(node, notes)
      type = if @todo_closures[node]
        @todo_closures[node]
      else
        #t = @typer.getInferredType(node).resolve
        #if t.kind_of? MethodType
          parent_type_future = @typer.getInferredType(node.parent)
          unless parent_type_future
            puts "#{CallSite(node.parent).name} call with block has no type at #{node.parent.position}"
            puts "  block type: #{@typer.getInferredType(node)}"
          end
          fs = CallFuture(parent_type_future).futures
          TypeFuture(fs.get(fs.size-1)).resolve
        #else
        #end
      end
      Map(notes).put node, type
    end
  end

  attr_accessor blockCloneMapOldNew:IdentityHashMap
  attr_accessor blockCloneMapNewOld:IdentityHashMap
  attr_accessor parent_scope_to_binding_name:Map

  def finish
    closures = []
    scripts = ArrayList.new(@scripts)
    Collections.reverse(scripts)
    scripts.each do |s: Script|
      closures.addAll BlockFinder.new(@typer, @todo_closures).find(s).entrySet
    end

    closures_to_skip = []

    blockToBindings = LinkedHashMap.new # the list of bindings a block closes over
    bindingLocalNamesToTypes = LinkedHashMap.new

    bindingForBlocks = LinkedHashMap.new # the specific binding for a given block
    Collections.reverse(closures) # from outside to inside
    
    self.blockCloneMapOldNew = IdentityHashMap.new
    self.blockCloneMapNewOld = IdentityHashMap.new
    
    selff = self
    
    closures.each do |entry: Entry|
      block = Block(entry.getKey)
      block.whenCloned do |interim,new|
        old = selff.blockCloneMapNewOld.get(interim)
        selff.blockCloneMapNewOld.put(new,old)
#       blockCloneMapOldNew.remove(old)
        selff.blockCloneMapOldNew.put(old,new)
      end
      blockCloneMapOldNew.put(block,block)
      blockCloneMapNewOld.put(block,block)
    end

    self.parent_scope_to_binding_name = {}

    i = 0
    closures.each do |entry: Entry|
      @@log.fine "adjust bindings for block #{entry.getKey} #{entry.getValue} #{i}"
      i += 1
      uncloned_block = Block(entry.getKey)
      block = Block(blockCloneMapOldNew.get(uncloned_block))
      @@log.fine "#{typer.sourceContent block}"
      enclosing_node = find_enclosing_node(block)
      if enclosing_node.nil?
        # this likely means a macro exists and made things confusing 
        # by copying the tree
        @@log.fine "enclosing node was nil, removing  #{entry.getKey} #{entry.getValue} #{i}"
        closures_to_skip.add entry
        next
      end
      @@log.fine "enclosing node #{enclosing_node}"
      @@log.fine "#{typer.sourceContent enclosing_node}"

      ProxyCleanup.new.scan enclosing_node

      enclosing_b = get_body(enclosing_node)      
      if enclosing_b.nil?
        closures_to_skip.add entry
        next
      end
      bindingName = "b#{i}"
      bindingForBlocks.put uncloned_block, bindingName
      parent_scope = MirrorScope(get_scope(block)) 
      adjuster = BindingAdjuster.new(
        self,
        bindingName,
        parent_scope,
        blockToBindings,
        bindingLocalNamesToTypes)

      binding_generated = adjuster.adjust enclosing_b, block
      @@log.log(Level.FINE, "After adjusting: #{AstFormatter.new(Node(scripts.get(0)))}")
#    end
#
#    # ignore closures with no parents, they aren't in the final AST, maybe
#    closures.removeAll closures_to_skip
#
#    i = 0
#    closures.each do |entry: Entry|
#      @@log.fine "insert_closure #{entry.getKey} #{entry.getValue} #{i}"
#      i += 1

      if AstChecker.enabled
        scripts.each do |s: Script|
          s.accept(AstChecker.new,nil)
        end
      end
  
      block = Block(blockCloneMapOldNew.get(entry.getKey))
      parent_type = ResolvedType(entry.getValue)

      unless get_body(find_enclosing_node(block))
        @@log.fine "  enclosing node was nil, removing  #{entry.getKey} #{entry.getValue} #{i}"
        next
      end

      closure_name = temp_name_from_outer_scope(block, "Closure")
      closure_klass = build_class(block.position, parent_type, closure_name)

      # build closure class
      binding_list = Collection(blockToBindings.get(uncloned_block)) || Collections.emptyList
      binding_args = binding_list.map do |name: String|
        RequiredArgument.new(SimpleString.new(name), SimpleString.new(ResolvedType(bindingLocalNamesToTypes[name]).name))
      end

      args = Arguments.new(closure_klass.position,
                           binding_args,
                           Collections.emptyList,
                           nil,
                           Collections.emptyList,
                           nil)
      binding_assigns = binding_list.map do |name: String|
        FieldAssign.new(SimpleString.new(name), LocalAccess.new(SimpleString.new(name)), nil)
      end
      constructor = ConstructorDefinition.new(
        SimpleString.new('initialize'), args,
        SimpleString.new('void'), binding_assigns, nil)
      closure_klass.body.add(constructor)

      enclosing_b  = find_enclosing_body block
      if binding_assigns.isEmpty
        insert_into_body enclosing_b, closure_klass
      else 
        # insert after binding class, for happier typing
        # also causes weird issues in the compiler if things are out of order
        enclosing_b.insert(1, closure_klass)
      end

      block_scope = get_scope block
      if contains_methods(block)
        copy_methods(closure_klass, block, block_scope)
      else
        build_and_inject_methods(closure_klass, block, parent_type, block_scope)
      end

      closure_type = infer(closure_klass) # FIXME: this re-infers also the body of the method (which is the ex-body of the block), which is probably duplicate work.

      has_block_parent = block.findAncestor { |node| node.parent.kind_of?(Block) || node.parent.kind_of?(ClosureDefinition) } # block, or converted block

      binding_locals = binding_list.map do |name: String|
        if has_block_parent && !name.equals(parent_scope_to_binding_name[parent_scope])
          FieldAccess.new(SimpleString.new(name))
        else
          LocalAccess.new(SimpleString.new(name))
        end
      end
      target = makeTypeName(block.position, closure_type.resolve)
      new_node = Call.new(
        block.position, target,
        SimpleString.new("new"), 
        binding_locals, nil)
      @typer.workaroundASTBug new_node

      

      if block.parent.kind_of?(CallSite)
        parent = CallSite(block.parent)
        replace_block_with_closure_in_call parent, block, new_node
      else
        replace_synthetic_lambda_definiton_with_closure(SyntheticLambdaDefinition(block.parent),new_node)
      end

      @@log.fine "inferring new_node #{new_node}"
      infer new_node
      @@log.fine "inferring enclosing_b #{enclosing_b}"
      infer enclosing_b

      @@log.fine "done with #{enclosing_b}"
      @@log.log(Level.FINE, "Inferred AST:\n{0}", AstFormatter.new(enclosing_b))

      buf = java::io::ByteArrayOutputStream.new
      ps = java::io::PrintStream.new(buf)
      printer = TypePrinter2.new(@typer, ps)
      printer.scan(enclosing_b, nil)
      ps.close()
      @@log.fine("Inferred types for expr:\n#{String.new(buf.toByteArray)}")
    end
  end

  def add_todo(block: Block, parent_type: ResolvedType)
    return if parent_type.isError || block.parent.nil?
    
    rtype = BaseTypeFuture.new(block.position)
    rtype.resolved parent_type
  
    new_scope = @typer.addNestedScope block
    new_scope.selfType = rtype
    if contains_methods block
      @typer.infer block.body
    else
      @typer.inferClosureBlock block, method_for(parent_type)
    end

    script = block.findAncestor{|n| n.kind_of? Script}

    @todo_closures[block] = parent_type
    @scripts.add script
  end

  def insert_closure(block: Block, parent_type: ResolvedType)
    raise "BetterClosureBuilder doesn't support insert_closure"
  end

  def prepare_non_local_return_closure(block: Block, parent_type: ResolvedType): Node
    # generates closure classes, AND an exception type
    # and replaces the closure call with something like this:
    #
    # class MyNonLocalReturn < Throwable
    #   def initialize(return_value:`method return type`); @return_value = return_value; end
    #   def return_value; @return_value; end
    # end
    # begin
    #   call { raise MyNonLocalReturn, `value` }
    # rescue MyNonLocalReturn => e
    #   return e.return_value
    # end
    enclosing_node = find_enclosing_node block
    return_type = if enclosing_node.kind_of? MethodDefinition
                    methodType = infer(enclosing_node)
                    MethodFuture(methodType).returnType
                  elsif enclosing_node.kind_of? Script
                    future = AssignableTypeFuture.new block.position
                    future.assign(infer(enclosing_node), block.position)
                    future
                  end
    nlr_klass = define_nlr_exception block
    block = convert_returns_to_raises block, nlr_klass, return_type
    new_node = nlr_prepare block, parent_type, nlr_klass
    resolved = return_type.resolve

    raise "Unable to determine method return type before generating closure including non local return" unless resolved

    enclosing_body = get_body(enclosing_node)
    node_in_body = block.findAncestor { |node| node.parent.kind_of? NodeList }
    new_call = wrap_with_rescue block, nlr_klass, node_in_body, resolved
    node_in_body.parent.replaceChild node_in_body, new_call
    
    finish_nlr_exception block, nlr_klass, resolved
    insert_into_body enclosing_body, nlr_klass
    infer(nlr_klass)
    new_node
  end

  # creates closure class, inserts it
  # returns Call Node that is the instantiation of the closure
  def prepare(block: Block, parent_type: ResolvedType): Call
    Call(prepare_regular_closure(block, parent_type))
  end

  def prepare_regular_closure(block: Block, parent_type: ResolvedType): Node
    parent_scope = get_scope block
    klass = build_closure_class block, parent_type, parent_scope
    if contains_methods(block)
      copy_methods(klass, block, parent_scope)
    else
      build_and_inject_methods(klass, block, parent_type, parent_scope)
    end
    new_closure_call_node(block, klass)
  end

  def replace_block_with_closure_in_call(parent: CallSite, block: Block, new_node: Node): void
    if block == parent.block
      parent.block = nil
      parent.parameters.add(new_node)
    else
      new_node.setParent(nil)
      parent.replaceChild(block, new_node)
    end
  end
  
  def replace_synthetic_lambda_definiton_with_closure(parent: SyntheticLambdaDefinition, new_node: Node): void
    parentparent = parent.parent
    new_node.setParent(nil)
    if parentparent.kind_of?(CallSite) && CallSite(parentparent).target!=parent # then the SyntheticLambdaDefinition is not a child of the CallSite itself, but (most likely?) a child of its arguments. FIXME: It is weird that the parent of a child of X is not X. 
      CallSite(parentparent).parameters.replaceChild(parent,new_node)
    else
      parentparent.replaceChild(parent,new_node)
    end
  end

  def find_enclosing_body block: Block
    enclosing_node = find_enclosing_node block
    get_body enclosing_node
  end

  def find_enclosing_method_body block: Block
    enclosing_node = find_enclosing_method block
    get_body enclosing_node
  end

  def get_body node: Node
    # TODO create an interface for nodes with bodies
    if node.kind_of?(MethodDefinition)
      MethodDefinition(node).body
    elsif node.kind_of?(Script)
      Script(node).body
    elsif node.kind_of?(Block)
      Block(node).body
    else
      raise "Unknown type for finding a body #{node.getClass}"
    end
  end

  def find_enclosing_node block: Node
    if block.parent
{'fraction_non_alphanumeric': 0.08314606741573034, 'fraction_numerical': 0.0012295950816196735, 'mean_word_length': 3.091153512575889, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7138263', 'n_tokens_mistral': 14181, 'n_tokens_neox': 13727, 'n_words': 3511}