entrypoint decoding & packing improved

pytezos/encoding.py

@@ -1,4 +1,6 @@
 import base58
+import calendar
+from datetime import datetime
 
 
 def tb(l):
@@ -263,6 +265,13 @@ def parse_address(data: bytes):
         return base58_encode(data[1:], tz_prefixes[b'\x00' + data[:1]]).decode()
 
 
+def parse_contract(data: bytes):
+    res = parse_address(data[:22])
+    if len(data) > 22:
+        res += f'%{data[22:].decode()}'
+    return res
+
+
 def forge_bool(value) -> bytes:
     return b'\xff' if value else b'\x00'
 
@@ -273,3 +282,18 @@ def forge_array(data, len_bytes=4) -> bytes:
 
 def forge_base58(value) -> bytes:
     return base58_decode(value.encode())
+
+
+def forge_timestamp(value) -> int:
+    assert isinstance(value, str)
+    dt = datetime.strptime(value, '%Y-%m-%dT%H:%M:%SZ')
+    return calendar.timegm(dt.utctimetuple())
+
+
+def forge_contract(value) -> bytes:
+    parts = value.split('%')
+    address, entrypoint = (parts[0], parts[1]) if len(parts) == 2 else (parts[0], 'default')
+    res = forge_address(address)
+    if entrypoint != 'default':
+        res += entrypoint.encode()
+    return res

pytezos/michelson/forge.py

@@ -1,4 +1,5 @@
-from pytezos.encoding import forge_array
+from pytezos.encoding import forge_array, forge_address, forge_public_key, forge_timestamp, forge_contract, \
+    forge_base58
 
 prim_tags = {
     'parameter': b'\x00',
@@ -223,5 +224,60 @@ def forge_script(script):
     return forge_array(code) + forge_array(storage)
 
 
-def pack(data):
+def prepack_micheline(val_expr, type_expr):
+    def convert(val_node, type_node):
+        prim = type_node['prim']
+        is_string = isinstance(val_node, dict) and val_node.get('string')
+
+        if prim in ['set', 'list']:
+            return [
+                prepack_micheline(item_node, type_node['args'][0])
+                for item_node in val_node
+            ]
+        elif prim in ['map', 'big_map']:
+            return [
+                {'prim': 'Elt',
+                 'args': [prepack_micheline(elt_node['args'][i], type_node['args'][i]) for i in [0, 1]]}
+                for elt_node in val_node
+            ]
+        elif prim == 'pair':
+            return [
+                {'prim': 'Pair',
+                 'args': [prepack_micheline(val_node['args'][i], type_node['args'][i]) for i in [0, 1]]}
+            ]
+        elif prim == 'option':
+            if val_node['prim'] == 'Some':
+                return {'prim': 'Some',
+                        'args': [prepack_micheline(val_node['args'][0], type_node['args'][0])]}
+            else:
+                return val_node
+        elif prim == 'or':
+            idx = {'Left': 0, 'Right': 1}[val_node['prim']]
+            return {'prim': val_node['prim'],
+                    'args': [prepack_micheline(val_node['args'][0], type_node['args'][idx])]}
+        elif prim == 'lambda':
+            # TODO: PUSH, SELF, CONTRACT
+            return val_node
+        elif prim == 'chain_id' and is_string:
+            return {'bytes': forge_base58(val_node['string']).hex()}
+        elif prim == 'signature' and is_string:
+            return {'bytes': forge_base58(val_node['string']).hex()}
+        elif prim == 'key_hash' and is_string:
+            return {'bytes': forge_address(val_node['string'], tz_only=True).hex()}
+        elif prim == 'key' and is_string:
+            return {'bytes': forge_public_key(val_node['string']).hex()}
+        elif prim == 'address' and is_string:
+            return {'bytes': forge_address(val_node['string']).hex()}
+        elif prim == 'contract' and is_string:
+            return {'bytes': forge_contract(val_node['string']).hex()}
+        elif prim == 'timestamp' and is_string:
+            return {'int': forge_timestamp(val_node['string'])}
+        else:
+            return val_node
+
+    return convert(val_expr, type_expr)
+
+
+def pack(val_expr, type_expr):
+    data = prepack_micheline(val_expr, type_expr)
     return b'\x05' + forge_micheline(data)

pytezos/repl/blockchain.py

@@ -5,10 +5,11 @@
 from pytezos.repl.context import Context
 from pytezos.repl.types import assert_stack_type, Mutez, ChainID, Address, Contract, Option, assert_equal_types, \
     KeyHash, Timestamp, expr_equal
-from pytezos.repl.parser import MichelsonTypeCheckError
+from pytezos.repl.parser import get_entry_expr
 
 MAINNET_CHAIN_ID = 'NetXdQprcVkpaWU'
 UNIT_TYPE_EXPR = {'prim': 'unit'}
+SELF_ADDRESS = 'KT1VG2WtYdSWz5E7chTeAdDPZNy2MpP8pTfL'
 
 
 def get_network_by_chain_id(chain_id: ChainID):
@@ -57,10 +58,13 @@ def do_chain_id(ctx: Context, prim, args, annots):
 def do_self(ctx: Context, prim, args, annots):
     p_type_expr = ctx.get('parameter')
     assert p_type_expr, f'parameter type is not initialized'
-    address = ctx.get('ADDRESS')
-    assert address, f'ADDRESS is not initialized'
-    res = Contract.new(str(address), type_expr=p_type_expr)
-    ctx.push(res, annots=annots)
+
+    entrypoint = next((a for a in annots if a[0] == '%'), '%default')
+    ctx.print(f' use {entrypoint};')
+
+    p_type_expr = get_entry_expr(p_type_expr, entrypoint)
+    res = Contract.new(SELF_ADDRESS + entrypoint, type_expr=p_type_expr)
+    ctx.push(res, annots=[a for a in annots if a[0] != '%'])
 
 
 @instruction('SENDER')

pytezos/repl/helpers.py

@@ -4,7 +4,7 @@
 from pytezos.encoding import is_kt
 from pytezos.repl.control import instruction, do_interpret
 from pytezos.repl.context import Context, StackItem
-from pytezos.repl.parser import get_int, get_string, parse_prim_expr
+from pytezos.repl.parser import get_int, get_string, parse_prim_expr, get_entry_expr
 from pytezos.repl.types import Pair, Mutez, Address, ChainID, Timestamp
 
 helpers_prim = ['DUMP', 'PRINT', 'DROP_ALL', 'EXPAND', 'RUN', 'PATCH', 'UNSET', 'INCLUDE']
@@ -35,6 +35,11 @@ def do_expand(ctx: Context, prim, args, annots):
 def do_run(ctx: Context, prim, args, annots):
     p_type_expr = ctx.get('parameter')
     assert p_type_expr, f'parameter type is not initialized'
+
+    entrypoint = next((a for a in annots if a[0] == '%'), '%default')
+    ctx.print(f' use {entrypoint};')
+
+    p_type_expr = get_entry_expr(p_type_expr, entrypoint)
     parameter = StackItem.parse(args[0], p_type_expr)
 
     s_type_expr = ctx.get('storage')

pytezos/repl/parser.py

@@ -1,6 +1,4 @@
 import functools
-import calendar
-from datetime import datetime
 from typing import Tuple, Callable
 
 from pytezos.michelson.converter import micheline_to_michelson
@@ -113,6 +111,25 @@ def get_bytes(val_expr):
     return bytes.fromhex(get_core_val(val_expr, core_type='bytes'))
 
 
+def get_entry_expr(expr, field_annot):
+    def _get(node):
+        assert_type(node, dict)
+        if field_annot in node.get('annots', []):
+            return node
+
+        for arg in node.get('args', []):
+            res = _get(arg)
+            if res:
+                return res
+
+    entry = _get(expr)
+    if not entry and field_annot == '%default':
+        entry = expr
+
+    assert entry, (expr, field_annot)
+    return entry
+
+
 def expr_equal(a, b):
     if type(a) != type(b):
         return False
@@ -294,14 +311,9 @@ def parse_big_map(val_expr, type_args):
     return parse_map(val_expr, type_args)
 
 
-def strptime(s) -> int:
-    dt = datetime.strptime(s, '%Y-%m-%dT%H:%M:%SZ')
-    return calendar.timegm(dt.utctimetuple())
-
-
 @primitive('timestamp')
 def parse_timestamp(val_expr, type_args):
-    return dispatch_core_map(val_expr, {'int': int, 'string': strptime})
+    return dispatch_core_map(val_expr, {'int': int, 'string': encoding.forge_timestamp})
 
 
 @primitive('mutez')
@@ -318,8 +330,11 @@ def parse_address(val_expr, type_args):
 
 
 @primitive('contract', args_len=1)
-def parse_contract(val_expr, type_args):  # TODO: entrypoint?
-    return parse_address(val_expr, type_args)
+def parse_contract(val_expr, type_args):
+    return dispatch_core_map(val_expr, {
+        'bytes': lambda x: encoding.parse_contract(bytes.fromhex(x)),
+        'string': lambda x: x
+    })
 
 
 @primitive('operation')

pytezos/repl/structures.py

@@ -111,7 +111,7 @@ def do_none(ctx: Context, prim, args, annots):
 @instruction('PACK')
 def do_pack(ctx: Context, prim, args, annots):
     top = ctx.pop1()
-    res = Bytes(pack(top.val_expr))
+    res = Bytes(pack(top.val_expr, top.type_expr))
     ctx.push(res, annots=annots)
 
 

pytezos/repl/types.py

@@ -376,10 +376,12 @@ def new(cls, address):
 class Contract(StackItem, prim='contract', args_len=1):
 
     @classmethod
-    def new(cls, address, type_expr):
-        assert is_pkh(address) or is_kt(address), f'expected address, got {address}'
-        return cls(val=address,
-                   val_expr={'string': address},
+    def new(cls, contract, type_expr):
+        assert is_pkh(contract[:36]) or is_kt(contract[:36]), f'expected contract, got {contract}'
+        if len(contract) > 36:
+            assert contract[36] == '%', f'expected contract, got {contract}'
+        return cls(val=contract,
+                   val_expr={'string': contract},
                    type_expr={'prim': cls.prim, 'args': [type_expr]})
 
 
@@ -438,7 +440,7 @@ def partial_apply(self, item: StackItem):
         push = {'prim': 'PUSH', 'args': [item.type_expr, item.val_expr]}
         pair = {'prim': 'PAIR'}
         code = self.val_expr['args'][0]
-        assert_type(code, list), f'expected instruction sequence'
+        assert_type(code, list), f'expected instruction sequence, got {code}'
         return type(self)(
             val_expr={'prim': 'Lambda', 'args': [[push, pair] + code]},
             type_expr={'prim': 'lambda', 'args': [r_type_expr, self.type_expr['args'][1]]}

tests/opcodes/contracts/car.tz

@@ -1,3 +1,3 @@
-parameter (pair nat nat);
+parameter (pair (nat :l) (nat :r));
 storage nat;
 code { CAR; CAR ; NIL operation ; PAIR }

tests/opcodes/contracts/cdr.tz

@@ -1,3 +1,3 @@
-parameter (pair nat nat);
+parameter (pair (nat :l) (nat :r));
 storage nat;
 code { CAR; CDR ; NIL operation ; PAIR }

tests/opcodes/contracts/check_signature.tz

@@ -1,8 +1,10 @@
 parameter key;
 storage (pair signature string);
-code { DUP; DUP;
+code {
+       DUP; DUP;
        DIP{ CDR; DUP; CAR;
-            DIP{CDR; PACK ; BLAKE2B}; PAIR};
-       CAR; DIP {UNPAIR}; CHECK_SIGNATURE;
+            DIP{CDR; PACK}};
+       CAR; CHECK_SIGNATURE;
        IF {} {FAIL} ;
        CDR; NIL operation ; PAIR};
+

tests/opcodes/contracts/create_contract.tz

@@ -11,4 +11,4 @@ code { DROP;
              { CDR;
                NIL operation;
                PAIR; } };
-       DROP; SOME; NIL operation; PAIR} # Ending calling convention stuff
+       DIP {SOME;NIL operation};CONS ; PAIR} # Ending calling convention stuff

tests/opcodes/contracts/self_with_default_entrypoint.tz

@@ -0,0 +1,19 @@
+parameter (or (or (nat %A) (bool %B)) (or %maybe_C (unit %default) (string %C)));
+storage unit;
+code {
+       DROP;
+       SELF; DROP;
+       # Refers to entrypoint A of the current contract.
+       SELF %A; DROP;
+       # Refers to the default entry of the current contract
+       SELF %default; PACK;
+       # "SELF" w/o annotation also refers to the default
+       # entry of the current contract. Internally, they are equal.
+       SELF; PACK; ASSERT_CMPEQ;
+       # The following instruction would not typecheck:
+       #    SELF %D,
+       # since there is no entrypoint D.
+       UNIT;
+       NIL operation;
+       PAIR;
+     }

tests/opcodes/contracts/self_with_entrypoint.tz

@@ -0,0 +1,26 @@
+parameter (or (or (nat %A) (bool %B)) (or %maybe_C (unit %Z) (string %C)));
+storage unit;
+code {
+       DROP;
+       # Refers to entrypoint A of the current contract.
+       SELF %A; PACK @Apacked;
+       # Refers to the default entry of the current contract
+       SELF %default; PACK @defpacked; DUP; DIP { SWAP }; ASSERT_CMPNEQ;
+       # "SELF" w/o annotation also refers to the default
+       # entry of the current contract
+       SELF; PACK @selfpacked; ASSERT_CMPEQ;
+
+       # Verify the types of the different entrypoints.  CAST is noop
+       # if its argument is convertible with the type of the top of
+       # the stack. is conver
+       SELF %A; CAST (contract nat); DROP;
+       SELF %B; CAST (contract bool); DROP;
+       SELF %maybe_C; CAST (contract (or (unit) (string))); DROP;
+       SELF %Z; CAST (contract unit); DROP;
+       SELF; CAST (contract (or (or (nat %A) (bool %B)) (or %maybe_C (unit %Z) (string %C)))); DROP;
+       SELF %default; CAST (contract (or (or (nat %A) (bool %B)) (or %maybe_C (unit %Z) (string %C)))); DROP;
+
+       UNIT;
+       NIL operation;
+       PAIR;
+     }