Implement basic support for instruction offsets

smalidea/src/main/java/org/jf/smalidea/SmaliTokens.java

@@ -34,6 +34,7 @@ package org.jf.smalidea;
 import com.google.common.collect.Maps;
 import com.intellij.openapi.editor.colors.TextAttributesKey;
 import com.intellij.psi.tree.IElementType;
+import com.intellij.psi.tree.TokenSet;
 import org.jf.smali.smaliParser;
 
 import java.lang.reflect.Field;
@@ -162,6 +163,8 @@ public class SmaliTokens {
     @SuppressWarnings({"UnusedDeclaration"}) public static IElementType VOID_TYPE;
     @SuppressWarnings({"UnusedDeclaration"}) public static IElementType VTABLE_INDEX;
 
+    public static final TokenSet INSTRUCTION_TOKENS;
+
     static {
         Map<String, TextAttributesKey> tokenColors = Maps.newHashMap();
 
@@ -309,5 +312,55 @@ public class SmaliTokens {
                 throw new RuntimeException(ex);
             }
         }
+
+        INSTRUCTION_TOKENS = TokenSet.create(
+                INSTRUCTION_FORMAT10t,
+                INSTRUCTION_FORMAT10x,
+                INSTRUCTION_FORMAT10x_ODEX,
+                INSTRUCTION_FORMAT11n,
+                INSTRUCTION_FORMAT11x,
+                INSTRUCTION_FORMAT12x_OR_ID,
+                INSTRUCTION_FORMAT12x,
+                INSTRUCTION_FORMAT20bc,
+                INSTRUCTION_FORMAT20t,
+                INSTRUCTION_FORMAT21c_FIELD,
+                INSTRUCTION_FORMAT21c_FIELD_ODEX,
+                INSTRUCTION_FORMAT21c_STRING,
+                INSTRUCTION_FORMAT21c_TYPE,
+                INSTRUCTION_FORMAT21ih,
+                INSTRUCTION_FORMAT21lh,
+                INSTRUCTION_FORMAT21s,
+                INSTRUCTION_FORMAT21t,
+                INSTRUCTION_FORMAT22b,
+                INSTRUCTION_FORMAT22c_FIELD,
+                INSTRUCTION_FORMAT22c_FIELD_ODEX,
+                INSTRUCTION_FORMAT22c_TYPE,
+                INSTRUCTION_FORMAT22cs_FIELD,
+                INSTRUCTION_FORMAT22s_OR_ID,
+                INSTRUCTION_FORMAT22s,
+                INSTRUCTION_FORMAT22t,
+                INSTRUCTION_FORMAT22x,
+                INSTRUCTION_FORMAT23x,
+                INSTRUCTION_FORMAT30t,
+                INSTRUCTION_FORMAT31c,
+                INSTRUCTION_FORMAT31i_OR_ID,
+                INSTRUCTION_FORMAT31i,
+                INSTRUCTION_FORMAT31t,
+                INSTRUCTION_FORMAT32x,
+                INSTRUCTION_FORMAT35c_METHOD,
+                INSTRUCTION_FORMAT35c_METHOD_ODEX,
+                INSTRUCTION_FORMAT35c_TYPE,
+                INSTRUCTION_FORMAT35mi_METHOD,
+                INSTRUCTION_FORMAT35ms_METHOD,
+                INSTRUCTION_FORMAT3rc_METHOD,
+                INSTRUCTION_FORMAT3rc_METHOD_ODEX,
+                INSTRUCTION_FORMAT3rc_TYPE,
+                INSTRUCTION_FORMAT3rmi_METHOD,
+                INSTRUCTION_FORMAT3rms_METHOD,
+                INSTRUCTION_FORMAT51l,
+                ARRAY_DATA_DIRECTIVE,
+                PACKED_SWITCH_DIRECTIVE,
+                SPARSE_SWITCH_DIRECTIVE
+        );
     }
 }

smalidea/src/main/java/org/jf/smalidea/psi/impl/SmaliInstruction.java

@@ -31,10 +31,21 @@
 
 package org.jf.smalidea.psi.impl;
 
+import com.intellij.lang.ASTNode;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+import org.jf.dexlib2.Opcode;
+import org.jf.dexlib2.Opcodes;
+import org.jf.smalidea.SmaliTokens;
 import org.jf.smalidea.psi.SmaliCompositeElementFactory;
 import org.jf.smalidea.psi.SmaliElementTypes;
 
 public class SmaliInstruction extends SmaliCompositeElement {
+    private static final int NO_OFFSET = -1;
+
+    @Nullable private Opcode opcode;
+    private int offset = NO_OFFSET;
+
     public static final SmaliCompositeElementFactory FACTORY = new SmaliCompositeElementFactory() {
         @Override public SmaliCompositeElement createElement() {
             return new SmaliInstruction();
@@ -44,4 +55,30 @@ public class SmaliInstruction extends SmaliCompositeElement {
     public SmaliInstruction() {
         super(SmaliElementTypes.INSTRUCTION);
     }
+
+    @NotNull public Opcode getOpcode() {
+        if (opcode == null) {
+            ASTNode instructionNode = findChildByType(SmaliTokens.INSTRUCTION_TOKENS);
+            // this should be impossible, based on the parser definition
+            assert instructionNode != null;
+
+            // TODO: put a project level Opcodes instance with the appropriate api level somewhere
+            opcode = new Opcodes(15).getOpcodeByName(instructionNode.getText());
+            assert opcode != null;
+        }
+        return opcode;
+    }
+
+    public int getOffset() {
+        if (offset == NO_OFFSET) {
+            SmaliInstruction previousInstruction = findPrevSiblingByClass(SmaliInstruction.class);
+            if (previousInstruction == null) {
+                offset = 0;
+            } else {
+                // TODO: handle variable size instructions
+                offset = previousInstruction.getOffset() + previousInstruction.getOpcode().format.size;
+            }
+        }
+        return offset;
+    }
 }

smalidea/src/main/java/org/jf/smalidea/psi/impl/SmaliMethod.java

@@ -31,6 +31,7 @@
 
 package org.jf.smalidea.psi.impl;
 
+import com.intellij.debugger.SourcePosition;
 import com.intellij.lang.ASTNode;
 import com.intellij.psi.*;
 import com.intellij.psi.PsiModifier.ModifierConstant;
@@ -105,6 +106,19 @@ public class SmaliMethod extends SmaliStubBasedPsiElement<SmaliMethodStub>
         return null;
     }
 
+    @NotNull public List<SmaliInstruction> getInstructions() {
+        return findChildrenByType(SmaliElementTypes.INSTRUCTION);
+    }
+
+    @Nullable public SourcePosition getSourcePositionForCodeOffset(int offset) {
+        for (SmaliInstruction instruction: getInstructions()) {
+            if (instruction.getOffset() >= offset) {
+                return SourcePosition.createFromElement(instruction);
+            }
+        }
+        return null;
+    }
+
     public int getRegisterCount() {
         SmaliRegistersStatement registersStatement = findChildByClass(SmaliRegistersStatement.class);
         if (registersStatement == null) {

smalidea/src/main/java/org/jf/smalidea/util/PsiUtils.java

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2014, Google Inc.
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *     * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *     * Neither the name of Google Inc. nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+package org.jf.smalidea.util;
+
+import com.intellij.psi.PsiElement;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+public class PsiUtils {
+    @Nullable
+    public static <T> T findPrevSiblingByClass(@Nonnull PsiElement element, @Nonnull Class<T> cls) {
+        PsiElement prev = element.getPrevSibling();
+        while (true) {
+            if (prev == null) {
+                return null;
+            } else if (cls.isInstance(prev)) {
+                return (T)prev;
+            }
+            prev = prev.getPrevSibling();
+        }
+    }
+}

smalidea/src/test/java/org/jf/smalidea/SmaliInstructionTest.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright 2014, Google Inc.
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *     * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *     * Neither the name of Google Inc. nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+package org.jf.smalidea;
+
+import com.intellij.psi.PsiElement;
+import com.intellij.testFramework.fixtures.LightCodeInsightFixtureTestCase;
+import org.jf.dexlib2.Opcode;
+import org.jf.smalidea.psi.impl.SmaliFile;
+import org.jf.smalidea.psi.impl.SmaliInstruction;
+import org.junit.Assert;
+
+public class SmaliInstructionTest extends LightCodeInsightFixtureTestCase {
+    public void testSingleInstruction() {
+        String text =
+                ".class public Lmy/pkg/blah; .super Ljava/lang/Object;\n" +
+                        ".method blah(IJLjava/lang/String;)V\n" +
+                        "    .locals 0\n" +
+                        "    r<ref>eturn-void\n" +
+                        ".end method";
+
+        SmaliFile file = (SmaliFile)myFixture.addFileToProject("my/pkg/blah.smali",
+                text.replace("<ref>", ""));
+
+        PsiElement leafElement = file.findElementAt(text.indexOf("<ref>"));
+        Assert.assertNotNull(leafElement);
+        SmaliInstruction instructionElement = (SmaliInstruction)leafElement.getParent();
+        Assert.assertNotNull(instructionElement);
+
+        Assert.assertEquals(Opcode.RETURN_VOID, instructionElement.getOpcode());
+        Assert.assertEquals(0, instructionElement.getOffset());
+    }
+
+    public void testMultipleInstructions() {
+        String text =
+                ".class public Lmy/pkg/blah; .super Ljava/lang/Object;\n" +
+                        ".method blah(IJLjava/lang/String;)I\n" +
+                        "    .locals 1\n" +
+                        "    const v0, 1234\n" +
+                        "    r<ref>eturn v0\n" +
+                        ".end method";
+
+        SmaliFile file = (SmaliFile)myFixture.addFileToProject("my/pkg/blah.smali",
+                text.replace("<ref>", ""));
+
+        PsiElement leafElement = file.findElementAt(text.indexOf("<ref>"));
+        Assert.assertNotNull(leafElement);
+        SmaliInstruction instructionElement = (SmaliInstruction)leafElement.getParent();
+        Assert.assertNotNull(instructionElement);
+
+        Assert.assertEquals(Opcode.RETURN, instructionElement.getOpcode());
+        Assert.assertEquals(6, instructionElement.getOffset());
+    }
+}

smalidea/src/test/java/org/jf/smalidea/SmaliMethodTest.java

@@ -31,11 +31,15 @@
 
 package org.jf.smalidea;
 
+import com.intellij.debugger.SourcePosition;
 import com.intellij.psi.PsiElement;
 import com.intellij.testFramework.fixtures.LightCodeInsightFixtureTestCase;
 import junit.framework.Assert;
+import org.jf.dexlib2.Opcode;
 import org.jf.smalidea.psi.impl.*;
 
+import java.util.List;
+
 public class SmaliMethodTest extends LightCodeInsightFixtureTestCase {
     public void testMethodRegisters() {
         String text =
@@ -183,4 +187,91 @@ public class SmaliMethodTest extends LightCodeInsightFixtureTestCase {
         Assert.assertFalse(smaliMethod.getParameterList().getParameters()[1].isVarArgs());
         Assert.assertFalse(smaliMethod.getParameterList().getParameters()[2].isVarArgs());
     }
+
+    private static final String instructionsTestClass =
+            ".class public Lmy/pkg/blah; .super Ljava/lang/Object;\n" +
+                    ".method public getRandomParentType(I)I\n" +
+                    "    .registers 4\n" +
+                    "    .param p1, \"edge\"    # I\n" +
+                    "\n" +
+                    "    .prologue\n" +
+                    "    const/4 v1, 0x2\n" +
+                    "\n" +
+                    "    .line 179\n" +
+                    "    if-nez p1, :cond_5\n" +
+                    "\n" +
+                    "    move v0, v1\n" +
+                    "\n" +
+                    "    .line 185\n" +
+                    "    :goto_4\n" +
+                    "    return v0\n" +
+                    "\n" +
+                    "    .line 182\n" +
+                    "    :cond_5\n" +
+                    "    if-ne p1, v1, :cond_f\n" +
+                    "\n" +
+                    "    .line 183\n" +
+                    "    sget-object v0, Lorg/jf/Penroser/PenroserApp;->random:Ljava/util/Random;\n" +
+                    "\n" +
+                    "    const/4 v1, 0x3\n" +
+                    "\n" +
+                    "    invoke-virtual {v0, v1}, Ljava/util/Random;->nextInt(I)I\n" +
+                    "\n" +
+                    "    move-result v0\n" +
+                    "\n" +
+                    "    goto :goto_4\n" +
+                    "\n" +
+                    "    .line 185\n" +
+                    "    :cond_f\n" +
+                    "    sget-object v0, Lorg/jf/Penroser/PenroserApp;->random:Ljava/util/Random;\n" +
+                    "\n" +
+                    "    invoke-virtual {v0, v1}, Ljava/util/Random;->nextInt(I)I\n" +
+                    "\n" +
+                    "    move-result v0\n" +
+                    "\n" +
+                    "    goto :goto_4\n" +
+                    ".end method";
+
+    public void testGetInstructions() {
+        String text = instructionsTestClass;
+
+        SmaliFile file = (SmaliFile)myFixture.addFileToProject("my/pkg/blah.smali", text);
+        SmaliClass smaliClass = file.getPsiClass();
+        SmaliMethod smaliMethod = smaliClass.getMethods()[0];
+
+        List<SmaliInstruction> instructions = smaliMethod.getInstructions();
+        Assert.assertEquals(14, instructions.size());
+    }
+
+    private void checkSourcePosition(SmaliMethod smaliMethod, int codeOffset, Opcode opcode) {
+        SourcePosition sourcePosition = smaliMethod.getSourcePositionForCodeOffset(codeOffset);
+        Assert.assertNotNull(sourcePosition);
+
+        SmaliInstruction instruction = (SmaliInstruction)sourcePosition.getElementAt();
+        Assert.assertEquals(opcode, instruction.getOpcode());
+        Assert.assertEquals(codeOffset, instruction.getOffset());
+    }
+
+    public void testGetSourcePositionForCodeOffset() {
+        String text = instructionsTestClass;
+
+        SmaliFile file = (SmaliFile)myFixture.addFileToProject("my/pkg/blah.smali", text);
+        SmaliClass smaliClass = file.getPsiClass();
+        SmaliMethod smaliMethod = smaliClass.getMethods()[0];
+
+        checkSourcePosition(smaliMethod, 0, Opcode.CONST_4);
+        checkSourcePosition(smaliMethod, 2, Opcode.IF_NEZ);
+        checkSourcePosition(smaliMethod, 6, Opcode.MOVE);
+        checkSourcePosition(smaliMethod, 8, Opcode.RETURN);
+        checkSourcePosition(smaliMethod, 10, Opcode.IF_NE);
+        checkSourcePosition(smaliMethod, 14, Opcode.SGET_OBJECT);
+        checkSourcePosition(smaliMethod, 18, Opcode.CONST_4);
+        checkSourcePosition(smaliMethod, 20, Opcode.INVOKE_VIRTUAL);
+        checkSourcePosition(smaliMethod, 26, Opcode.MOVE_RESULT);
+        checkSourcePosition(smaliMethod, 28, Opcode.GOTO);
+        checkSourcePosition(smaliMethod, 30, Opcode.SGET_OBJECT);
+        checkSourcePosition(smaliMethod, 34, Opcode.INVOKE_VIRTUAL);
+        checkSourcePosition(smaliMethod, 40, Opcode.MOVE_RESULT);
+        checkSourcePosition(smaliMethod, 42, Opcode.GOTO);
+    }
 }