Add support for invoke-custom and related structures in the parser and lexer

9 files changed, 91 insertions, 15 deletions

diff --git a/smali/src/main/antlr/smaliParser.g b/smali/src/main/antlr/smaliParser.g
index 2d5eccaa..abb5a521 100644
--- a/smali/src/main/antlr/smaliParser.g
+++ b/smali/src/main/antlr/smaliParser.g
@@ -41,6 +41,7 @@ tokens {
   ARRAY_DATA_DIRECTIVE;
   ARRAY_TYPE_PREFIX;
   ARROW;
+  AT;
   BOOL_LITERAL;
   BYTE_LITERAL;
   CATCH_DIRECTIVE;
@@ -106,11 +107,13 @@ tokens {
   INSTRUCTION_FORMAT31i_OR_ID;
   INSTRUCTION_FORMAT31t;
   INSTRUCTION_FORMAT32x;
+  INSTRUCTION_FORMAT35c_CALL_SITE;
   INSTRUCTION_FORMAT35c_METHOD;
   INSTRUCTION_FORMAT35c_METHOD_ODEX;
   INSTRUCTION_FORMAT35c_TYPE;
   INSTRUCTION_FORMAT35mi_METHOD;
   INSTRUCTION_FORMAT35ms_METHOD;
+  INSTRUCTION_FORMAT3rc_CALL_SITE;
   INSTRUCTION_FORMAT3rc_METHOD;
   INSTRUCTION_FORMAT3rc_METHOD_ODEX;
   INSTRUCTION_FORMAT3rc_TYPE;
@@ -125,6 +128,8 @@ tokens {
   LOCALS_DIRECTIVE;
   LONG_LITERAL;
   METHOD_DIRECTIVE;
+  METHOD_HANDLE_TYPE_FIELD;
+  METHOD_HANDLE_TYPE_METHOD;
   MEMBER_NAME;
   NEGATIVE_INTEGER_LITERAL;
   NULL_LITERAL;
@@ -176,6 +181,7 @@ tokens {
   I_ANNOTATION;
   I_ANNOTATION_ELEMENT;
   I_SUBANNOTATION;
+  I_ENCODED_METHOD_HANDLE;
   I_ENCODED_FIELD;
   I_ENCODED_METHOD;
   I_ENCODED_ENUM;
@@ -224,8 +230,10 @@ tokens {
   I_STATEMENT_FORMAT31i;
   I_STATEMENT_FORMAT31t;
   I_STATEMENT_FORMAT32x;
+  I_STATEMENT_FORMAT35c_CALL_SITE;
   I_STATEMENT_FORMAT35c_METHOD;
   I_STATEMENT_FORMAT35c_TYPE;
+  I_STATEMENT_FORMAT3rc_CALL_SITE;
   I_STATEMENT_FORMAT3rc_METHOD;
   I_STATEMENT_FORMAT3rc_TYPE;
   I_STATEMENT_FORMAT45cc_METHOD;
@@ -236,12 +244,13 @@ tokens {
   I_STATEMENT_SPARSE_SWITCH;
   I_REGISTER_RANGE;
   I_REGISTER_LIST;
+  I_CALL_SITE_EXTRA_ARGUMENTS;
+  I_CALL_SITE_REFERENCE;
 }
 
 @header {
 package org.jf.smali;
 
-import org.jf.dexlib2.Format;
 import org.jf.dexlib2.Opcode;
 import org.jf.dexlib2.Opcodes;
 }
@@ -549,6 +558,8 @@ simple_name
   | PRIMITIVE_TYPE -> SIMPLE_NAME[$PRIMITIVE_TYPE]
   | VOID_TYPE -> SIMPLE_NAME[$VOID_TYPE]
   | ANNOTATION_VISIBILITY -> SIMPLE_NAME[$ANNOTATION_VISIBILITY]
+  | METHOD_HANDLE_TYPE_FIELD
+  | METHOD_HANDLE_TYPE_METHOD
   | INSTRUCTION_FORMAT10t -> SIMPLE_NAME[$INSTRUCTION_FORMAT10t]
   | INSTRUCTION_FORMAT10x -> SIMPLE_NAME[$INSTRUCTION_FORMAT10x]
   | INSTRUCTION_FORMAT10x_ODEX -> SIMPLE_NAME[$INSTRUCTION_FORMAT10x_ODEX]
@@ -568,6 +579,7 @@ simple_name
   | INSTRUCTION_FORMAT23x -> SIMPLE_NAME[$INSTRUCTION_FORMAT23x]
   | INSTRUCTION_FORMAT31i_OR_ID -> SIMPLE_NAME[$INSTRUCTION_FORMAT31i_OR_ID]
   | INSTRUCTION_FORMAT31t -> SIMPLE_NAME[$INSTRUCTION_FORMAT31t]
+  | INSTRUCTION_FORMAT35c_CALL_SITE -> SIMPLE_NAME[$INSTRUCTION_FORMAT35c_CALL_SITE]
   | INSTRUCTION_FORMAT35c_METHOD -> SIMPLE_NAME[$INSTRUCTION_FORMAT35c_METHOD]
   | INSTRUCTION_FORMAT35c_METHOD_ODEX -> SIMPLE_NAME[$INSTRUCTION_FORMAT35c_METHOD_ODEX]
   | INSTRUCTION_FORMAT35c_TYPE -> SIMPLE_NAME[$INSTRUCTION_FORMAT35c_TYPE]
@@ -636,7 +648,9 @@ literal
   | array_literal
   | subannotation
   | type_field_method_literal
-  | enum_literal;
+  | enum_literal
+  | method_handle_literal
+  | method_prototype;
 
 parsed_integer_literal returns[int value]
   : integer_literal { $value = LiteralTools.parseInt($integer_literal.text); };
@@ -699,6 +713,19 @@ type_field_method_literal
   | PRIMITIVE_TYPE
   | VOID_TYPE;
 
+call_site_reference
+  : simple_name OPEN_PAREN STRING_LITERAL COMMA method_prototype (COMMA literal)* CLOSE_PAREN AT method_reference
+    -> ^(I_CALL_SITE_REFERENCE simple_name STRING_LITERAL method_prototype ^(I_CALL_SITE_EXTRA_ARGUMENTS literal*)
+         method_reference);
+
+method_handle_reference
+  : METHOD_HANDLE_TYPE_FIELD AT field_reference -> METHOD_HANDLE_TYPE_FIELD field_reference
+  | METHOD_HANDLE_TYPE_METHOD AT method_reference -> METHOD_HANDLE_TYPE_METHOD method_reference;
+
+method_handle_literal
+  : method_handle_reference
+  -> ^(I_ENCODED_METHOD_HANDLE method_handle_reference);
+
 method_reference
   : (reference_type_descriptor ARROW)? member_name method_prototype
   -> reference_type_descriptor? member_name method_prototype;
@@ -796,8 +823,6 @@ instruction_format31i
   : INSTRUCTION_FORMAT31i
   | INSTRUCTION_FORMAT31i_OR_ID -> INSTRUCTION_FORMAT31i[$INSTRUCTION_FORMAT31i_OR_ID];
 
-
-
 instruction
   : insn_format10t
   | insn_format10x
@@ -829,11 +854,13 @@ instruction
   | insn_format31i
   | insn_format31t
   | insn_format32x
+  | insn_format35c_call_site
   | insn_format35c_method
   | insn_format35c_type
   | insn_format35c_method_odex
   | insn_format35mi_method
   | insn_format35ms_method
+  | insn_format3rc_call_site
   | insn_format3rc_method
   | insn_format3rc_method_odex
   | insn_format3rc_type
@@ -1016,6 +1043,12 @@ insn_format32x
     INSTRUCTION_FORMAT32x REGISTER COMMA REGISTER
     -> ^(I_STATEMENT_FORMAT32x[$start, "I_STATEMENT_FORMAT32x"] INSTRUCTION_FORMAT32x REGISTER REGISTER);
 
+insn_format35c_call_site
+  : //e.g. invoke-custom {v0, v1}, call_site_name
+    // OR invoke-custom {v0, v1}, {"doSomething", (LCustom;Ljava/lang/String;)Ljava/lang/String;, "just testing"}, BootstrapLinker;->normalLink(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;Ljava/lang/Object;)Ljava/lang/invoke/CallSite;
+    INSTRUCTION_FORMAT35c_CALL_SITE OPEN_BRACE register_list CLOSE_BRACE COMMA call_site_reference
+    -> ^(I_STATEMENT_FORMAT35c_CALL_SITE[$start, "I_STATEMENT_FORMAT35c_CALL_SITE"] INSTRUCTION_FORMAT35c_CALL_SITE register_list call_site_reference);
+
 insn_format35c_method
   : //e.g. invoke-virtual {v0,v1} java/io/PrintStream/print(Ljava/lang/Stream;)V
     INSTRUCTION_FORMAT35c_METHOD OPEN_BRACE register_list CLOSE_BRACE COMMA method_reference
@@ -1047,6 +1080,12 @@ insn_format35ms_method
       throwOdexedInstructionException(input, $INSTRUCTION_FORMAT35ms_METHOD.text);
     };
 
+insn_format3rc_call_site
+  : //e.g. invoke-custom/range {v0 .. v1}, call_site_name
+    // OR invoke-custom/range {v0 .. v1}, {"doSomething", (LCustom;Ljava/lang/String;)Ljava/lang/String;, "just testing"}, BootstrapLinker;->normalLink(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;Ljava/lang/Object;)Ljava/lang/invoke/CallSite;
+    INSTRUCTION_FORMAT3rc_CALL_SITE OPEN_BRACE register_range CLOSE_BRACE COMMA call_site_reference
+    -> ^(I_STATEMENT_FORMAT3rc_CALL_SITE[$start, "I_STATEMENT_FORMAT3rc_CALL_SITE"] INSTRUCTION_FORMAT3rc_CALL_SITE register_range call_site_reference);
+
 insn_format3rc_method
   : //e.g. invoke-virtual/range {v25..v26}, java/lang/StringBuilder/append(Ljava/lang/String;)Ljava/lang/StringBuilder;
     INSTRUCTION_FORMAT3rc_METHOD OPEN_BRACE register_range CLOSE_BRACE COMMA method_reference
diff --git a/smali/src/main/jflex/smaliLexer.jflex b/smali/src/main/jflex/smaliLexer.jflex
index 2cf3c869..70378ec5 100644
--- a/smali/src/main/jflex/smaliLexer.jflex
+++ b/smali/src/main/jflex/smaliLexer.jflex
@@ -408,6 +408,14 @@ Type = {PrimitiveType} | {ClassDescriptor} | {ArrayPrefix} ({ClassDescriptor} |
     "vtable@0x" {HexDigit}+ { return newToken(VTABLE_INDEX); }
     "field@0x" {HexDigit}+ { return newToken(FIELD_OFFSET); }
 
+    "instance-get" | "instance-put" | "static-get" | "static-put" {
+        return newToken(METHOD_HANDLE_TYPE_FIELD);
+    }
+
+    "instance-invoke" | "static-invoke" {
+        return newToken(METHOD_HANDLE_TYPE_METHOD);
+    }
+
     # [^\r\n]* { return newToken(LINE_COMMENT, true); }
 }
 
@@ -567,6 +575,10 @@ Type = {PrimitiveType} | {ClassDescriptor} | {ArrayPrefix} ({ClassDescriptor} |
         return newToken(INSTRUCTION_FORMAT32x);
     }
 
+    "invoke-custom" {
+        return newToken(INSTRUCTION_FORMAT35c_CALL_SITE);
+    }
+
     "invoke-virtual" | "invoke-super" | "invoke-direct" | "invoke-static" | "invoke-interface" {
         return newToken(INSTRUCTION_FORMAT35c_METHOD);
     }
@@ -587,6 +599,10 @@ Type = {PrimitiveType} | {ClassDescriptor} | {ArrayPrefix} ({ClassDescriptor} |
         return newToken(INSTRUCTION_FORMAT35ms_METHOD);
     }
 
+    "invoke-custom/range" {
+        return newToken(INSTRUCTION_FORMAT3rc_CALL_SITE);
+    }
+
     "invoke-virtual/range" | "invoke-super/range" | "invoke-direct/range" | "invoke-static/range" |
     "invoke-interface/range" {
         return newToken(INSTRUCTION_FORMAT3rc_METHOD);
@@ -668,6 +684,7 @@ Type = {PrimitiveType} | {ClassDescriptor} | {ArrayPrefix} ({ClassDescriptor} |
     "}" { return newToken(CLOSE_BRACE); }
     "(" { return newToken(OPEN_PAREN); }
     ")" { return newToken(CLOSE_PAREN); }
+    "@" { return newToken(AT); }
     [\r\n\t ]+ { return newToken(WHITE_SPACE, true); }
     <<EOF>> { return newToken(EOF); }
 }
diff --git a/smali/src/test/resources/LexerTest/DirectiveTest.tokens b/smali/src/test/resources/LexerTest/DirectiveTest.tokens
index 06f4da73..9b822202 100644
--- a/smali/src/test/resources/LexerTest/DirectiveTest.tokens
+++ b/smali/src/test/resources/LexerTest/DirectiveTest.tokens
@@ -53,7 +53,7 @@ INVALID_TOKEN(".end blah")
 INVALID_TOKEN(".local1234")
 INVALID_TOKEN(".super1234")
 SUPER_DIRECTIVE(".super")
-INVALID_TOKEN("@")
+AT("@")
 SUPER_DIRECTIVE(".super")
 SUPER_DIRECTIVE(".super")
 INVALID_TOKEN(".supeer")
diff --git a/smali/src/test/resources/LexerTest/InstructionTest.smali b/smali/src/test/resources/LexerTest/InstructionTest.smali
index 62e51006..829a3e37 100644
--- a/smali/src/test/resources/LexerTest/InstructionTest.smali
+++ b/smali/src/test/resources/LexerTest/InstructionTest.smali
@@ -209,6 +209,7 @@ sparse-switch
 move/16
 move-wide/16
 move-object/16
+invoke-custom
 invoke-virtual
 invoke-super
 invoke-direct
@@ -221,6 +222,7 @@ throw-verification-error
 execute-inline
 invoke-virtual-quick
 invoke-super-quick
+invoke-custom/range
 invoke-virtual/range
 invoke-super/range
 invoke-direct/range
diff --git a/smali/src/test/resources/LexerTest/InstructionTest.tokens b/smali/src/test/resources/LexerTest/InstructionTest.tokens
index dfbd584b..efab4ec4 100644
--- a/smali/src/test/resources/LexerTest/InstructionTest.tokens
+++ b/smali/src/test/resources/LexerTest/InstructionTest.tokens
@@ -209,6 +209,7 @@ INSTRUCTION_FORMAT31t("sparse-switch")
 INSTRUCTION_FORMAT32x("move/16")
 INSTRUCTION_FORMAT32x("move-wide/16")
 INSTRUCTION_FORMAT32x("move-object/16")
+INSTRUCTION_FORMAT35c_CALL_SITE("invoke-custom")
 INSTRUCTION_FORMAT35c_METHOD("invoke-virtual")
 INSTRUCTION_FORMAT35c_METHOD("invoke-super")
 INSTRUCTION_FORMAT35c_METHOD("invoke-direct")
@@ -221,6 +222,7 @@ INSTRUCTION_FORMAT20bc("throw-verification-error")
 INSTRUCTION_FORMAT35mi_METHOD("execute-inline")
 INSTRUCTION_FORMAT35ms_METHOD("invoke-virtual-quick")
 INSTRUCTION_FORMAT35ms_METHOD("invoke-super-quick")
+INSTRUCTION_FORMAT3rc_CALL_SITE("invoke-custom/range")
 INSTRUCTION_FORMAT3rc_METHOD("invoke-virtual/range")
 INSTRUCTION_FORMAT3rc_METHOD("invoke-super/range")
 INSTRUCTION_FORMAT3rc_METHOD("invoke-direct/range")
diff --git a/smali/src/test/resources/LexerTest/MiscTest.smali b/smali/src/test/resources/LexerTest/MiscTest.smali
index 32bbd272..b9dfea5e 100644
--- a/smali/src/test/resources/LexerTest/MiscTest.smali
+++ b/smali/src/test/resources/LexerTest/MiscTest.smali
@@ -45,6 +45,13 @@ illegal-method-access
 class-change-error
 instantiation-error
 
+instance-invoke
+static-invoke
+instance-get
+instance-put
+static-get
+static-put
+
 inline@0xABCD
 inline@0x0123
 inline@0x0123ABCD
diff --git a/smali/src/test/resources/LexerTest/MiscTest.tokens b/smali/src/test/resources/LexerTest/MiscTest.tokens
index ed5142d8..7ba62a7c 100644
--- a/smali/src/test/resources/LexerTest/MiscTest.tokens
+++ b/smali/src/test/resources/LexerTest/MiscTest.tokens
@@ -45,6 +45,13 @@ VERIFICATION_ERROR_TYPE("illegal-method-access")
 VERIFICATION_ERROR_TYPE("class-change-error")
 VERIFICATION_ERROR_TYPE("instantiation-error")
 
+METHOD_HANDLE_TYPE_METHOD("instance-invoke")
+METHOD_HANDLE_TYPE_METHOD("static-invoke")
+METHOD_HANDLE_TYPE_FIELD("instance-get")
+METHOD_HANDLE_TYPE_FIELD("instance-put")
+METHOD_HANDLE_TYPE_FIELD("static-get")
+METHOD_HANDLE_TYPE_FIELD("static-put")
+
 INLINE_INDEX("inline@0xABCD")
 INLINE_INDEX("inline@0x0123")
 INLINE_INDEX("inline@0x0123ABCD")
@@ -57,15 +64,15 @@ FIELD_OFFSET("field@0xABCD")
 FIELD_OFFSET("field@0x0123")
 FIELD_OFFSET("field@0x0123ABCD")
 
-SIMPLE_NAME("inline") INVALID_TOKEN("@")
-SIMPLE_NAME("inline") INVALID_TOKEN("@") SIMPLE_NAME("zzz")
-SIMPLE_NAME("inline") INVALID_TOKEN("@") SIMPLE_NAME("abcd")
-SIMPLE_NAME("vtable") INVALID_TOKEN("@")
-SIMPLE_NAME("vtable") INVALID_TOKEN("@") SIMPLE_NAME("zzz")
-SIMPLE_NAME("vtable") INVALID_TOKEN("@") SIMPLE_NAME("abcd")
-SIMPLE_NAME("field") INVALID_TOKEN("@")
-SIMPLE_NAME("field") INVALID_TOKEN("@") SIMPLE_NAME("zzz")
-SIMPLE_NAME("field") INVALID_TOKEN("@") SIMPLE_NAME("abcd")
+SIMPLE_NAME("inline") AT("@")
+SIMPLE_NAME("inline") AT("@") SIMPLE_NAME("zzz")
+SIMPLE_NAME("inline") AT("@") SIMPLE_NAME("abcd")
+SIMPLE_NAME("vtable") AT("@")
+SIMPLE_NAME("vtable") AT("@") SIMPLE_NAME("zzz")
+SIMPLE_NAME("vtable") AT("@") SIMPLE_NAME("abcd")
+SIMPLE_NAME("field") AT("@")
+SIMPLE_NAME("field") AT("@") SIMPLE_NAME("zzz")
+SIMPLE_NAME("field") AT("@") SIMPLE_NAME("abcd")
 
 INVALID_TOKEN("+") POSITIVE_INTEGER_LITERAL("0")
 INVALID_TOKEN("+") POSITIVE_INTEGER_LITERAL("10")
diff --git a/smali/src/test/resources/LexerTest/SymbolTest.smali b/smali/src/test/resources/LexerTest/SymbolTest.smali
index 6b11061b..d443a652 100644
--- a/smali/src/test/resources/LexerTest/SymbolTest.smali
+++ b/smali/src/test/resources/LexerTest/SymbolTest.smali
@@ -9,6 +9,7 @@
 }
 (
 )
+@
  
 	
 	 
diff --git a/smali/src/test/resources/LexerTest/SymbolTest.tokens b/smali/src/test/resources/LexerTest/SymbolTest.tokens
index 009c9324..9d967f77 100644
--- a/smali/src/test/resources/LexerTest/SymbolTest.tokens
+++ b/smali/src/test/resources/LexerTest/SymbolTest.tokens
@@ -8,5 +8,6 @@ OPEN_BRACE("{") WHITE_SPACE(" ") CLOSE_BRACE("}") WHITE_SPACE(" ") OPEN_PAREN("(
 OPEN_BRACE("{") WHITE_SPACE("\n")
 CLOSE_BRACE("}") WHITE_SPACE("\n")
 OPEN_PAREN("(") WHITE_SPACE("\n")
-CLOSE_PAREN(")")
+CLOSE_PAREN(")") WHITE_SPACE("\n")
+AT("@")
 WHITE_SPACE("\n \n\t\n\t \n\t  \n\t \t\n \t\n  \t\n\r\r")
\ No newline at end of file