content: # Copyright 2015 Google Inc. All Rights Reserved.
  #
  # 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.
  
  import collections
  
  from .. import ir
  from .. import scalartypes as scalars
  from ..jvmops import *
  
  # Copy propagation - when one register is moved to another, keep track and replace
  # all loads with loads from the original register (as long as it hasn't since been
  # overwritten). Note that stores won't be removed, since they may still be needed
  # in some cases, but if they are unused, they'll be removed in a subsequent pass
  # As usual, assume no iincs
  
  # A set of registers that currently are copies of each other.
  class _CopySet:
      def __init__(self, key):
          self.root = key
          self.set = {key}
          self.q = [] # keep track of insertion order in case root is overwritten
  
      def add(self, key):
          assert self.set
          self.set.add(key)
          self.q.append(key)
  
      def remove(self, key):
          self.set.remove(key)
          # Heuristic - use oldest element still in set as new root
          while self.q and self.root not in self.set:
              self.root = self.q.pop(0)
  
      def copy(self):
          new = _CopySet(self.root)
          new.set = self.set.copy()
          new.q = self.q[:]
          return new
  
  # Map registers to CopySets
  class _CopySetsMap:
      def __init__(self):
          self.lookup = {}
  
      def _get(self, key): return self.lookup.setdefault(key, _CopySet(key))
  
      def clobber(self, key):
          self._get(key).remove(key)
          del self.lookup[key]
  
      def move(self, dest, src):
          # return false if the corresponding instructions should be removed
          s_set = self._get(src)
          d_set = self._get(dest)
          if s_set is d_set:
              # src and dest are copies of same value, so we can remove
              return False
          d_set.remove(dest)
          s_set.add(dest)
          self.lookup[dest] = s_set
          return True
  
      def load(self, key):
          return self._get(key).root
  
      def copy(self):
          copies = {}
          new = _CopySetsMap()
          for k, v in self.lookup.items():
              if v not in copies:
                  copies[v] = v.copy()
              new.lookup[k] = copies[v]
          return new
  
  def copyPropagation(irdata):
      instrs = irdata.flat_instructions
      replace = {}
  
      single_pred_infos = {}
  
      prev = None
      current = _CopySetsMap()
      for instr in instrs:
          # reset all info when control flow is merged
          if instr in irdata.jump_targets:
              # try to use info if this was a single predecessor forward jump
              if prev and not prev.fallsthrough() and irdata.target_pred_counts.get(instr) == 1:
                  current = single_pred_infos.get(instr, _CopySetsMap())
              else:
                  current = _CopySetsMap()
  
          elif isinstance(instr, ir.RegAccess):
              key = instr.key
              if instr.store:
                  # check if previous instr was a load
                  if isinstance(prev, ir.RegAccess) and not prev.store:
                      if not current.move(dest=key, src=prev.key):
                          replace[prev] = []
                          replace[instr] = []
                  else:
                      current.clobber(key)
              else:
                  root_key = current.load(key)
                  if key != root_key:
                      assert instr not in replace
                      # replace with load from root register instead
                      replace[instr] = [ir.RegAccess(root_key[0], root_key[1], False)]
  
          else:
              for target in instr.targets():
                  label = irdata.labels[target]
                  if irdata.target_pred_counts.get(label) == 1:
                      single_pred_infos[label] = current.copy()
  
          prev = instr
      irdata.replaceInstrs(replace)
  
  def _isRemoveable(instr):
      # can remove if load or const since we know there are no side effects
      # note - instr may be None
      if isinstance(instr, ir.RegAccess) and not instr.store:
          return True
      return isinstance(instr, (ir.PrimConstant, ir.OtherConstant))
  
  def removeUnusedRegisters(irdata):
      # Remove stores to registers that are not read from anywhere in the method
      instrs = irdata.flat_instructions
      used = set()
      for instr in instrs:
          if isinstance(instr, ir.RegAccess) and not instr.store:
              used.add(instr.key)
  
      replace = {}
      prev = None
      for instr in instrs:
          if isinstance(instr, ir.RegAccess) and instr.key not in used:
              assert instr.store
              # if prev instruction is load or const, just remove it and the store
              # otherwise, replace the store with a pop
              if _isRemoveable(prev):
                  replace[prev] = []
                  replace[instr] = []
              else:
                  replace[instr] = [ir.Pop2() if instr.wide else ir.Pop()]
          prev = instr
      irdata.replaceInstrs(replace)
  
  # Allocate registers to JVM registers on a first come, first serve basis
  # For simplicity, parameter registers are preserved as is
  def simpleAllocateRegisters(irdata):
      instrs = irdata.flat_instructions
      regmap = {v:i for i,v in enumerate(irdata.initial_args)}
      nextreg = len(irdata.initial_args)
  
      for instr in instrs:
          if isinstance(instr, ir.RegAccess):
              if instr.key not in regmap:
                  regmap[instr.key] = nextreg
                  nextreg += 1
                  if instr.wide:
                      nextreg += 1
              instr.calcBytecode(regmap[instr.key])
      irdata.numregs = nextreg
  
  # Sort registers by number of uses so that more frequently used registers will
  # end up in slots 0-3 or 4-255 and benefit from the shorter instruction forms
  # For simplicity, parameter registers are still preserved as is with one exception
  def sortAllocateRegisters(irdata):
      instrs = irdata.flat_instructions
  
      use_counts = collections.Counter()
      for instr in instrs:
          if isinstance(instr, ir.RegAccess):
              use_counts[instr.key] += 1
  
      regs = irdata.initial_args[:]
      rest = sorted(use_counts, key=lambda k:(-use_counts[k], k))
      for key in rest:
          # If key is a param, it was already added at the beginning
          if key not in irdata.initial_args:
              regs.append(key)
              if scalars.iswide(key[1]):
                  regs.append(None)
  
      # Sometimes the non-param regsisters are used more times than the param registers
      # and it is beneificial to swap them (which requires inserting code at the
      # beginning of the method to move the value if the param is not unused)
      # This is very complicated to do in general, so the following code only does
      # this in one specific circumstance which should nevertheless be sufficient
      # to capture the majority of the benefit
      # Specificially, it only swaps at most one register, and only in the case that
      # it is nonwide and there is a nonwide parameter in the first 4 slots that
      # it can be swapped with. Also, it doesn't bother to check if param is unused.
      candidate_i = max(4, len(irdata.initial_args))
      # make sure candidate is valid, nonwide register
      if len(regs) > candidate_i and regs[candidate_i] is not None:
          candidate = regs[candidate_i]
          if not scalars.iswide(candidate[1]) and use_counts[candidate] >= 3:
              for i in range(min(4, len(irdata.initial_args))):
                  # make sure target is not wide
                  if regs[i] is None or regs[i+1] is None:
                      continue
  
                  target = regs[i]
                  if use_counts[candidate] > use_counts[target] + 3:
                      # swap register assignments
                      regs[i], regs[candidate_i] = candidate, target
                      # add move instructions at beginning of method
                      load = ir.RegAccess.raw(i, target[1], False)
                      store = ir.RegAccess(target[0], target[1], True)
                      instrs = [load, store] + instrs
                      irdata.flat_instructions = instrs
                      break
  
      # Now generate bytecode from the selected register allocations
      irdata.numregs = len(regs)
      regmap = {v:i for i,v in enumerate(regs) if v is not None}
      for instr in instrs:
          if instr.bytecode is None and isinstance(instr, ir.RegAccess):
              instr.calcBytecode(regmap[instr.key])
  