Implement "entity local" feature - variables attached to a specific entity

Author: simon816

asm_reader.py

@@ -60,6 +60,11 @@ def __next__(self):
             const = self.next_constant()
             self.end_of_line()
             return const
+        elif char == '@':
+            self.skip(1)
+            ent_local = self.next_entity_local()
+            self.end_of_line()
+            return ent_local
         elif char == '_':
             self.skip(1)
             return self.next_local_label()
@@ -106,6 +111,10 @@ def next_constant(self):
         value = self.read_ref()
         return ('const', (name, value))
 
+    def next_entity_local(self):
+        name = self.read_symbol()
+        return ('entity_local', name)
+
     def next_local_label(self):
         return ('local_label', self.read_label())
 

assembler.py

@@ -45,6 +45,8 @@ def parse(self, text, filename=''):
             if token == 'const':
                 name, ref = arg
                 self.define_const(name, self.resolve_ref(*ref))
+            elif token == 'entity_local':
+                self.define_entity_local(arg)
             elif token == 'label':
                 self.handle_label(arg)
             elif token == 'instruction':
@@ -70,6 +72,9 @@ def define_const(self, name, value):
             raise RuntimeError('Constant %r already defined' % name)
         self.constants[name] = value
 
+    def define_entity_local(self, name):
+        self.define_const(name, EntityLocal(name))
+
     def get_const(self, name):
         return self.constants[name]
 
@@ -180,11 +185,11 @@ def add_cast(self, type, cmd, dest):
         }[type]
         # store result of cmd into nbt of the specified type
         self.add_command(ExecuteChain() \
-                         .store('result').entity(data_type).run(cmd))
+                         .store('result').entity(EntityTag, data_type).run(cmd))
         # move nbt value back out to scoreboard
         self.add_command(ExecuteChain() \
                          .store('result').score(EntityTag, dest) \
-                         .run(DataGet(data_type)))
+                         .run(DataGet(EntityTag, data_type)))
 
 
     def handle_directive(self, directive, value):

commands.py

@@ -49,7 +49,9 @@ def resolve(self, scope):
         pass
 
 class Ref(Resolvable):
-    pass
+
+    def selector(self, where):
+        return ETagSelector(where)
 
 class Var(Ref):
     def __init__(self, nameref, *args):
@@ -66,6 +68,16 @@ def __init__(self, loc):
     def resolve(self, scope):
         return scope.memory(self.loc)
 
+class EntityLocal(Ref):
+    def __init__(self, name):
+        self.name = name
+
+    def resolve(self, scope):
+        return scope.entity_local(self.name)
+
+    def selector(self, where):
+        return Selector('s', where)
+
 class SimpleResolve(Resolvable):
 
     def __init__(self, *args):
@@ -101,14 +113,29 @@ def str_pairs(items):
 
 class Selector(Resolvable):
 
-    def __init__(self, where):
+    def __init__(self, type, where):
+        self.type = type
         self.where = where
 
+    def resolve_params(self, scope):
+        if not self.where:
+            return {}
+        return self.where.resolve(scope)
+
     def resolve(self, scope):
-        where = {} if not self.where else self.where.resolve(scope)
-        return make_selector('e', tag=scope.entity_tag, limit=1, **where)
+        return make_selector(self.type, **self.resolve_params(scope))
+
+class ETagSelector(Selector):
+
+    def __init__(self, where):
+        super().__init__('e', where)
 
-EntityTag = Selector(None)
+    def resolve_params(self, scope):
+        where = super().resolve_params(scope)
+        where.update(tag=scope.entity_tag, limit=1)
+        return where
+
+EntityTag = ETagSelector(None)
 
 class Path(Resolvable):
 
@@ -161,6 +188,10 @@ def resolve(self, scope):
         return 'execute %s run %s' % (self.chain.resolve(scope),
                                       self.command.resolve(scope))
 
+def ensure_selector(sel_arg):
+    return sel_arg.as_selector() if isinstance(sel_arg, SelectorArgs) \
+           else sel_arg
+
 class ExecuteChain:
 
     def __init__(self):
@@ -188,7 +219,7 @@ def if_where(cond):
         return ExecuteChain().cond('if').where(cond)
 
     def where(self, select_arg):
-        return self.add('as', Selector(select_arg))
+        return self.add('as', ensure_selector(select_arg))
 
     def cond(self, cond_type):
         return ExecuteChain.Cond(self, cond_type)
@@ -203,7 +234,7 @@ def __init__(self, parent, cond_type):
             self.cond_type = cond_type
 
         def where(self, select_arg):
-            return self.add('entity', Selector(select_arg))
+            return self.add('entity', ensure_selector(select_arg))
 
         def score(self, target, t_objective, operator, source, s_objective):
             return self.add('score', target, t_objective, operator, source,
@@ -225,18 +256,20 @@ def __init__(self, parent, store_type):
             self.store_type = store_type
 
         def score(self, name, objective):
-            return self.add('score', name, objective)
+            return self.add('score', ensure_selector(name), objective)
 
-        def entity(self, data_type):
-            return self.add('entity', EntityTag, Path(data_type), data_type, 1)
+        def entity(self, target, data_type):
+            return self.add('entity', ensure_selector(target), \
+                            Path(data_type), data_type, 1)
 
 class DataGet(Command):
 
-    def __init__(self, path):
+    def __init__(self, target, path):
+        self.target = target
         self.path = Path(path)
 
     def resolve(self, scope):
-        return 'data get entity %s %s' % (EntityTag.resolve(scope),
+        return 'data get entity %s %s' % (self.target.resolve(scope),
                                           self.path.resolve(scope))
 
 class Function(Command):
@@ -290,7 +323,7 @@ def __init__(self, varref, value, where=None):
         assert self.allows_negative or value >= 0
         self.var = varref
         self.value = value
-        self.selector = Selector(where)
+        self.selector = varref.selector(where)
 
     def resolve(self, scope):
         return 'scoreboard players %s %s %s %d' % (
@@ -331,13 +364,13 @@ def __init__(self, left, right, where=None):
         assert isinstance(right, Ref)
         self.left = left
         self.right = right
-        self.selector = Selector(where)
+        self.left_sel = left.selector(where)
+        self.right_sel = right.selector(where)
 
     def resolve(self, scope):
-        selector = self.selector.resolve(scope)
         return 'scoreboard players operation %s %s %s %s %s' % (
-            selector, self.left.resolve(scope), self.op,
-            selector, self.right.resolve(scope))
+            self.left_sel.resolve(scope), self.left.resolve(scope), self.op,
+            self.right_sel.resolve(scope), self.right.resolve(scope))
 
 
 class OpAssign(Operation): op = '='
@@ -351,7 +384,16 @@ class OpIfGt(Operation): op = '>'
 class OpSwap(Operation): op = '><'
 
 class SelectorArgs(Resolvable):
-    pass
+
+    def as_selector(self):
+        return ETagSelector(self)
+
+class SimpleSelectorArgs(SelectorArgs):
+    def __init__(self, args):
+        self.args = args
+
+    def resolve(self, scope):
+        return dict(self.args)
 
 class SelRange(SelectorArgs):
     def __init__(self, varref, min=None, max=None):
@@ -361,6 +403,9 @@ def __init__(self, varref, min=None, max=None):
         self.min = min
         self.max = max
 
+    def as_selector(self):
+        return self.var.selector(self)
+
     def resolve(self, scope):
         name = self.var.resolve(scope)
         range = ''

compiler/asm_writer.py

@@ -25,6 +25,9 @@ def write_constant(self, name, value):
     def write_directive(self, name, value):
         self.write_line('#%s %s' % (name, value))
 
+    def write_entity_local(self, name):
+        self.write_line('@%s' % name)
+
     def write_subroutine(self, name):
         if name == '__setup__':
             return

compiler/compiler.py

@@ -20,6 +20,7 @@ def compile_program(self, program):
                              *(self.temp_registers + list(self.volatile_reg)))
         self.writer.write_instruction('MOV', '#0', '0')
         self.writer.write_instruction('MOV', '#0', '1')
+        self.entity_locals = {}
         self.optimizer = Optimizer(self)
         visitor = CompilerVisitor(self.optimizer.handle_insn)
         self.load_libs(visitor)
@@ -61,6 +62,10 @@ def handle_fn_end(self, insn):
     def handle_label(self, insn):
         self.writer.write_local_sub(insn.label)
 
+    def handle_entity_local(self, insn):
+        self.entity_locals[insn.offset] = insn.name
+        self.writer.write_entity_local(insn.name)
+
     def handle_jump(self, insn):
         self.writer.write_instruction('JMP',  self.label(insn.dest))
 
@@ -182,8 +187,8 @@ def handle_operation(self, insn):
             self.writer.write_instruction(opcode, right, tmp_ref)
             self._write_relative_move(tmp_ref, None, d_ref, d_off)
             self.free_slot(tmp_slot)
-        elif insn.op in [IR.Op.Eq, IR.Op.Neq, IR.Op.Lt,
-                         IR.Op.Gt, IR.Op.LtEq, IR.Op.GtEq]:
+        elif insn.op.__class__ in map(lambda o:o.__class__, [IR.Op.Eq,
+                    IR.Op.Neq, IR.Op.Lt, IR.Op.Gt, IR.Op.LtEq, IR.Op.GtEq]):
             jmpcode = {
                 IR.Op.Eq.__class__: 'JE',
                 IR.Op.Neq.__class__: 'JNE',
@@ -202,7 +207,17 @@ def handle_operation(self, insn):
             self._write_relative_move('#0', None, dest, dest_off)
             self.writer.write_local_sub(lbl)
         elif insn.op is IR.Op.LogAnd:
-            assert False, str(insn)
+            left = self.get_final_location(insn.left)
+            right = self.get_final_location(insn.right)
+            dest, dest_off = self.get_effective_location(insn.dest)
+            self._write_relative_move('#0', None, dest, dest_off)
+            self.writer.write_instruction('CMP', left, '#0')
+            lbl = 'and_end_%d' % abs(hash(insn)) # TODO
+            self.writer.write_instruction('JE', '_' + lbl)
+            self.writer.write_instruction('CMP', right, '#0')
+            self.writer.write_instruction('JE', '_' + lbl)
+            self._write_relative_move('#1', None, dest, dest_off)
+            self.writer.write_local_sub(lbl)
         elif insn.op is IR.Op.LogOr:
             assert False, str(insn)
         else:
@@ -323,6 +338,8 @@ def unpack(self, ref):
             return base, 0
         if isinstance(ref, IR.GlobalSlot):
             return None, ref.offset
+        if isinstance(ref, IR.EntityLocalSlot):
+            return self.entity_locals[ref.offset], None
         assert False, "Don't know how to handle " + str(ref)
 
     def __next_volatile(self):
@@ -414,12 +431,23 @@ def __str__(self):
     def relative_base(self):
         return IR.StackPointer
 
+class EntityLocalStorage(Storage):
+
+    def __str__(self):
+        return 'EntityLocal'
+
+    @property
+    def relative_base(self):
+        return IR.EntityLocalBase
+
 class ScopeManager:
 
-    def __init__(self):
+    def __init__(self, compiler):
+        self.compiler = compiler
         self.global_table = SymbolTable()
         self.current_table = self.global_table
         self.global_storage = GlobalStorage()
+        self.entity_local_storage = EntityLocalStorage()
         self.current_storage = self.global_storage
 
     def __enter__(self):
@@ -443,7 +471,13 @@ def lookup(self, name):
         return self.current_table.lookup(name)
 
     def declare_symbol(self, name, type, can_redeclare=False):
-        return self.current_table.declare(name, type, self.store(type),
+        if type is self.compiler.types.types['entity_local']:
+            # hijack entity_local
+            unit = self.entity_local_storage.insert(type)
+            self.compiler.emit(IR.EntityLocalVar(name, unit.offset))
+        else:
+            unit = self.store(type)
+        return self.current_table.declare(name, type, unit,
                                           can_redeclare)
 
     def store(self, type):
@@ -475,7 +509,9 @@ def __init__(self, on_instruction):
         super().__init__()
         self.on_instruction = on_instruction
         self.types = Types()
-        self.scope = ScopeManager()
+        self.types.add_type('entity_local', DecoratedType(IntType(),
+                    static=True, const=False))
+        self.scope = ScopeManager(self)
         self.current_function = None
         self.loop_attacher = LoopAttacher(None, None, self, None)
         self.python_functions = {}
@@ -489,6 +525,9 @@ def eliminate_offset(self, slot):
             if slot.base is IR.GlobalIndex:
                 assert isinstance(slot.offset, IR.LiteralInt), str(slot)
                 return IR.GlobalSlot(slot.offset.val, slot.type)
+            if slot.base is IR.EntityLocalBase:
+                assert isinstance(slot.offset, IR.LiteralInt), str(slot)
+                return IR.EntityLocalSlot(slot.offset.val, slot.type)
             base = self.eliminate_offset(slot.base)
             if isinstance(base, IR.GlobalSlot):
                 assert isinstance(slot.offset, IR.LiteralInt), str(slot)

compiler/ir.py

@@ -10,6 +10,7 @@ class IR:
     FunctionEnd = namedtuple('FunctionEnd', '')
 
     Label = namedtuple('Label', 'label')
+    EntityLocalVar = namedtuple('EntityLocalVar', 'name offset')
 
     Jump = namedtuple('Jump', 'dest')
     JumpIf = namedtuple('JumpIf', 'dest cond')
@@ -30,10 +31,12 @@ class IR:
     StackPointer = namedtuple('StackPointer', 'type')(IntType())
     BasePointer = namedtuple('BasePointer', 'type')(IntType())
     GlobalIndex = namedtuple('GlobalIndex', 'type')(IntType())
+    EntityLocalBase = namedtuple('EntityLocalBase', 'type')(IntType())
     ReturnRegister = namedtuple('ReturnRegister', 'type')(IntType())
 
     SlotOffset = namedtuple('SlotOffset', 'base offset type')
     GlobalSlot = namedtuple('GlobalSlot', 'offset type')
+    EntityLocalSlot = namedtuple('EntityLocalSlot', 'offset type')
     Dereference = namedtuple('Dereference', 'addr type')
     Free = namedtuple('Free', 'slot')
 
@@ -120,6 +123,7 @@ def __init__(self, downstream=None):
             IR.FunctionBegin: self.handle_fn_begin,
             IR.FunctionEnd: self.handle_fn_end,
             IR.Label: self.handle_label,
+            IR.EntityLocalVar: self.handle_entity_local,
             IR.Jump: self.handle_jump,
             IR.JumpIf: self.handle_jump_if,
             IR.JumpIfNot: self.handle_jump_if_not,
@@ -154,6 +158,9 @@ def handle_fn_end(self, insn):
     def handle_label(self, insn):
         self.emit(insn)
 
+    def handle_entity_local(self, insn):
+        self.emit(insn)
+
     def handle_jump(self, insn):
         self.emit(insn)