diff --git a/smali/build.gradle b/smali/build.gradle
index 4d558f5f..0877932d 100644
--- a/smali/build.gradle
+++ b/smali/build.gradle
@@ -77,12 +77,25 @@ dependencies {
     proguard depends.proguard
 }
 
-task generateAntlrSource(type: JavaExec) {
+task generateParserAntlrSource(type: JavaExec) {
     inputs.dir file(antlrSource)
     outputs.dir file(antlrOutput)
 
     mkdir(antlrOutput)
-    def grammars = fileTree(antlrSource).include('**/*.g')
+    def grammars = fileTree(antlrSource).include('**/smaliParser.g')
+
+    classpath = configurations.antlr3
+    main = 'org.antlr.Tool'
+    args '-fo', relativePath("${antlrOutput}/org/jf/smali")
+    args grammars.files
+}
+
+task generateTreeWalkerAntlrSource(type: JavaExec) {
+    inputs.dir file(antlrSource)
+    outputs.dir file(antlrOutput)
+
+    mkdir(antlrOutput)
+    def grammars = fileTree(antlrSource).include('**/smaliTreeWalker.g')
 
     classpath = configurations.antlr3
     main = 'org.antlr.Tool'
@@ -117,7 +130,7 @@ task generateJflexSource(type: JavaExec) {
     args grammars.files.join(' ')
 }
 
-compileJava.dependsOn generateAntlrSource, generateJflexSource
+compileJava.dependsOn generateParserAntlrSource, generateTreeWalkerAntlrSource, generateJflexSource
 compileTestJava.dependsOn generateTestAntlrSource
 
 // build a jar containing all dependencies