feat(lltz): added INDEX_ADDRESS and GET_ADDRESS_INDEX

lib/lltz_codegen/lltz_codegen.ml

@@ -185,6 +185,8 @@ let convert_primitive (prim : LLTZ.P.t) : Michelson.Ast.t =
   | Address -> address
   | Implicit_account -> implicit_account
   | Is_implicit_account -> is_implicit_account
+  | Index_address -> index_address
+  | Get_address_index -> get_address_index
   | Contract (annot, ty) -> contract ~annot (convert_type ty)
   | Pack -> pack
   | Unpack ty -> unpack (convert_type ty)

lib/lltz_ir/ast_builder.ml

@@ -585,6 +585,20 @@ module Default = struct
       (mk_type ~range (LLTZ.T.Option (mk_type ~range LLTZ.T.Key_hash)))
   ;;
 
+  let index_address ~range address =
+    create
+      ~range
+      (LLTZ.E.Prim (LLTZ.P.Index_address, [ address ]))
+      (mk_type ~range LLTZ.T.Nat)
+  ;;
+
+  let get_address_index ~range address =
+    create
+      ~range
+      (LLTZ.E.Prim (LLTZ.P.Get_address_index, [ address ]))
+      (mk_type ~range (LLTZ.T.Option (mk_type ~range LLTZ.T.Nat)))
+  ;;
+
   let contract ~range (opt, ty) address =
     create
       ~range
@@ -1242,6 +1256,8 @@ module With_dummy = struct
   let address contract = Default.address ~range:v contract
   let implicit_account key_hash = Default.implicit_account ~range:v key_hash
   let is_implicit_account address = Default.is_implicit_account ~range:v address
+  let index_address address = Default.index_address ~range:v address
+  let get_address_index address = Default.get_address_index ~range:v address
   let contract (opt, ty) address = Default.contract ~range:v (opt, ty) address
   let pack value = Default.pack ~range:v value
   let unpack ty value = Default.unpack ~range:v ty value

lib/lltz_ir/primitive.ml

@@ -44,6 +44,8 @@ type t =
   | Address
   | Implicit_account
   | Is_implicit_account
+  | Index_address
+  | Get_address_index
   | Contract of string option * Type.t
   | Pack
   | Unpack of Type.t

lib/lltz_michelson/lltz_michelson.ml

@@ -156,6 +156,8 @@ let convert_primitive (prim : LLTZ.P.t) : Michelson.Ast.t =
   | Address -> address
   | Implicit_account -> implicit_account
   | Is_implicit_account -> is_implicit_account
+  | Index_address -> index_address
+  | Get_address_index -> get_address_index
   | Contract (opt, ty) -> contract (convert_type ty) (* TODO: resolve tag option*)
   | Pack -> pack
   | Unpack ty -> unpack (convert_type ty)

lib/michelson/ast.ml

@@ -56,6 +56,8 @@ module Prim = struct
       | Create_contract
       | Implicit_account
       | Is_implicit_account
+      | Index_address
+      | Get_address_index
       | Dip
       | Drop
       | Dup
@@ -170,6 +172,8 @@ module Prim = struct
       | Create_contract -> "CREATE_CONTRACT"
       | Implicit_account -> "IMPLICIT_ACCOUNT"
       | Is_implicit_account -> "IS_IMPLICIT_ACCOUNT"
+      | Index_address -> "INDEX_ADDRESS"
+      | Get_address_index -> "GET_ADDRESS_INDEX"
       | Dip -> "DIP"
       | Drop -> "DROP"
       | Dup -> "DUP"
@@ -284,6 +288,8 @@ module Prim = struct
       | "CREATE_CONTRACT" -> Create_contract
       | "IMPLICIT_ACCOUNT" -> Implicit_account
       | "IS_IMPLICIT_ACCOUNT" -> Is_implicit_account
+      | "INDEX_ADDRESS" -> Index_address
+      | "GET_ADDRESS_INDEX" -> Get_address_index
       | "DIP" -> Dip
       | "DROP" -> Drop
       | "DUP" -> Dup
@@ -785,6 +791,8 @@ module Instruction = struct
   let if_none ~then_ ~else_ = prim ~arguments:[ seq then_; seq else_ ] (I If_none)
   let implicit_account = prim (I Implicit_account)
   let is_implicit_account = prim (I Is_implicit_account)
+  let index_address = prim (I Index_address)
+  let get_address_index = prim (I Get_address_index)
   let is_nat = prim (I Is_nat)
   let iter instrs = prim ~arguments:[ seq instrs ] (I Iter)
   let join_tickets = prim (I Join_tickets)

lib/michelson/optimisations/if_suffix_rewriter.ml

@@ -28,6 +28,7 @@ let is_injective : string -> bool = function
   | "CONS"
   | "IMPLICIT_ACCOUNT"
   | "IS_IMPLICIT_ACCOUNT"
+  | "GET_ADDRESS_INDEX" (* INDEX_ADDRESS is not pure ? *)
   | "EMPTY_MAP"
   | "EMPTY_SET"
   | "HASH_KEY"

lib/michelson/optimisations/michelson_base/primitive.ml

@@ -46,6 +46,8 @@ type 'ty prim1 =
   | Address
   | Implicit_account
   | Is_implicit_account
+  | Index_address
+  | Get_address_index
   | Contract of string option * 'ty
   | Pack
   | Unpack of 'ty

lib/michelson/optimisations/michelson_base/primitive.mli

@@ -47,6 +47,9 @@ type 'ty prim1 =
   | Address (** Address of a contract.*)
   | Implicit_account (** Implicit Account of a key_hash.*)
   | Is_implicit_account (** Is_implicit Account of an address.*)
+  | Index_address (** Index address and add to global cache if not already present *)
+  | Get_address_index
+  (** Get Some(index) for an address or None if not present in global cache *)
   | Contract of string option * 'ty (** Contract of an address and entrypoint.*)
   | Pack (** Packing values.*)
   | Unpack of 'ty (** Unpacking values.*)

lib/michelson/optimisations/michelson_base/typing.ml

@@ -108,6 +108,8 @@ let type_prim1 = function
   | Address
   | Implicit_account
   | Is_implicit_account
+  | Index_address
+  | Get_address_index
   | Pack
   | Set_delegate
   | Read_ticket

lib/michelson/optimisations/oasis_core/michelson.ml

@@ -1422,6 +1422,18 @@ let mi_is_implicit_account =
     | _ -> None)
 ;;
 
+let mi_index_address =
+  mk_spec_basic "INDEX_ADDRESS" ~arities:(1, 1) (function
+    | { mt = MT0 Address; _ } :: _ -> Some [ mt_nat ]
+    | _ -> None)
+;;
+
+let mi_get_address_index =
+  mk_spec_basic "GET_ADDRESS_INDEX" ~arities:(1, 1) (function
+    | { mt = MT0 Address; _ } :: _ -> Some [ mt_option mt_nat ]
+    | _ -> None)
+;;
+
 let mi_voting_power =
   mk_spec_basic "VOTING_POWER" ~arities:(1, 1) (function
     | { mt = MT0 Key_hash; _ } :: _ -> Some [ mt_nat ]
@@ -1523,6 +1535,8 @@ let spec_of_prim1 p =
   | Address -> mi_address
   | Implicit_account -> mi_implicit_account
   | Is_implicit_account -> mi_is_implicit_account
+  | Index_address -> mi_index_address
+  | Get_address_index -> mi_get_address_index
   | Voting_power -> mi_voting_power
   | Size -> mi_size
   | Car | Cdr -> assert false
@@ -1711,6 +1725,8 @@ let name_of_instr_exn = function
         | Address
         | Implicit_account
         | Is_implicit_account
+        | Index_address
+        | Get_address_index
         | Contract _
         | Pack
         | Unpack _
@@ -2134,6 +2150,8 @@ module Of_micheline = struct
          | "SELF_ADDRESS", [] -> MI0 Self_address
          | "IMPLICIT_ACCOUNT", [] -> MI1 Implicit_account
          | "IS_IMPLICIT_ACCOUNT", [] -> MI1 Is_implicit_account
+         | "INDEX_ADDRESS", [] -> MI1 Index_address
+         | "GET_ADDRESS_INDEX", [] -> MI1 Get_address_index
          | "TRANSFER_TOKENS", [] -> MI3 Transfer_tokens
          | "CHECK_SIGNATURE", [] -> MI3 Check_signature
          | "SET_DELEGATE", [] -> MI1 Set_delegate
@@ -2496,6 +2514,8 @@ module To_micheline = struct
           | Address
           | Implicit_account
           | Is_implicit_account
+          | Index_address
+          | Get_address_index
           | Pack
           | Hash_key
           | Blake2b

lib/michelson/optimisations/oasis_core/michelson_rewriter.ml

@@ -167,6 +167,8 @@ let rec may_fail = function
       | Address
       | Implicit_account
       | Is_implicit_account
+      | Index_address (* what about if the global cache fills up? *)
+      | Get_address_index
       | Pack
       | Hash_key
       | Blake2b

test/test_nodes.ml

@@ -2930,6 +2930,28 @@ let%expect_test "is_implicit_account address" =
     { PUSH address "tz1ABC123" ; IS_IMPLICIT_ACCOUNT } |}]
 ;;
 
+let%expect_test "index_address address" =
+  let e = index_address (address_const "tz1ABC123") in
+  test_expr e;
+  [%expect
+    {|
+    { PUSH address "tz1ABC123" ; INDEX_ADDRESS }
+
+    Optimised:
+    { PUSH address "tz1ABC123" ; INDEX_ADDRESS } |}]
+;;
+
+let%expect_test "get_address_index address" =
+  let e = get_address_index (address_const "tz1ABC123") in
+  test_expr e;
+  [%expect
+    {|
+    { PUSH address "tz1ABC123" ; GET_ADDRESS_INDEX }
+
+    Optimised:
+    { PUSH address "tz1ABC123" ; GET_ADDRESS_INDEX } |}]
+;;
+
 let%expect_test "contract opt (bool_ty) address" =
   let e = contract (None, bool_ty) (address_const "KT1XYZ") in
   test_expr e;