Uast: Rewrite visitor

plugins/lint/lint-api/src/com/android/tools/klint/client/api/LintDriver.java

@@ -1454,7 +1454,7 @@ public class LintDriver {
                         UastScanner scanner = (UastScanner) check;
                         UastVisitor customVisitor = scanner.createUastVisitor(context);
                         if (customVisitor != null) {
-                            UastChecker.INSTANCE.checkWithCustomHandler(
+                            UastChecker.INSTANCE.check(
                                     ideaProject, context.file, context, customVisitor);
                         } else {
                             UastChecker.INSTANCE.check(
@@ -1515,7 +1515,7 @@ public class LintDriver {
             for (UastScanner detector : detectors) {
                 UastVisitor customHandler = detector.createUastVisitor(context);
                 if (customHandler != null) {
-                    checker.checkWithCustomHandler(intellijProject, file, context, customHandler);
+                    checker.check(intellijProject, file, context, customHandler);
                 } else {
                     checker.check(intellijProject, file, detector, context);
                 }

plugins/lint/lint-checks/src/com/android/tools/klint/checks/CallSuperDetector.java

@@ -177,7 +177,7 @@ public class CallSuperDetector extends Detector implements UastScanner {
                 @NonNull UFunction methodDeclaration,
                 @NonNull UFunction superMethod) {
             SuperCallVisitor visitor = new SuperCallVisitor(context, superMethod);
-            visitor.process(methodDeclaration);
+            methodDeclaration.accept(visitor);
             return visitor.mCallsSuper;
         }
 
@@ -203,10 +203,8 @@ public class CallSuperDetector extends Detector implements UastScanner {
         }
 
         @Override
-        public void process(@NotNull UElement element) {
+        public boolean visitElement(@NotNull UElement node) {
-            if (!mCallsSuper) {
+            return mCallsSuper || super.visitElement(node);
-                super.process(element);
-            }
         }
     }
 }

plugins/lint/lint-checks/src/com/android/tools/klint/checks/CleanupDetector.java

@@ -269,7 +269,7 @@ public class CleanupDetector extends Detector implements UastScanner {
             }
         };
 
-        visitor.process(method);
+        method.accept(visitor);
         if (visitor.isCleanedUp() || visitor.variableEscapes()) {
             return;
         }
@@ -346,7 +346,7 @@ public class CleanupDetector extends Detector implements UastScanner {
                     }
                 };
 
-                commitVisitor.handle(method);
+                method.accept(commitVisitor);
                 if (commitVisitor.isCleanedUp() || commitVisitor.variableEscapes()) {
                     return true;
                 }
@@ -476,10 +476,8 @@ public class CleanupDetector extends Detector implements UastScanner {
         }
 
         @Override
-        public void process(@NotNull UElement element) {
+        public boolean visitElement(@NotNull UElement node) {
-            if (!mContainsCleanup) {
+            return mContainsCleanup || super.visitElement(node);
-                super.process(element);
-            }
         }
 
         protected abstract boolean isCleanupCall(@NonNull UCallExpression call);

plugins/lint/lint-checks/src/com/android/tools/klint/checks/CutPasteDetector.java

@@ -179,7 +179,7 @@ public class CutPasteDetector extends Detector implements UastScanner {
             @NonNull UCallExpression from,
             @NonNull UCallExpression to) {
         ReachableVisitor visitor = new ReachableVisitor(from, to);
-        visitor.process(method);
+        method.accept(visitor);
         return visitor.isReachable();
     }
 
@@ -214,16 +214,16 @@ public class CutPasteDetector extends Detector implements UastScanner {
             UExpression condition = node.getCondition();
             UExpression body = node.getThenBranch();
             UElement elseBody = node.getElseBranch();
-            process(condition);
+            condition.accept(this);
 
             if (body != null) {
                 boolean wasReachable = mReachable;
-                process(body);
+                body.accept(this);
                 mReachable = wasReachable;
             }
             if (elseBody != null) {
                 boolean wasReachable = mReachable;
-                process(elseBody);
+                elseBody.accept(this);
                 mReachable = wasReachable;
             }
 
@@ -231,10 +231,8 @@ public class CutPasteDetector extends Detector implements UastScanner {
         }
 
         @Override
-        public void process(@NotNull UElement element) {
+        public boolean visitElement(@NotNull UElement node) {
-            if (!mSeenEnd) {
+            return mSeenEnd;
-                super.process(element);
-            }
         }
     }
 }

plugins/lint/lint-checks/src/com/android/tools/klint/checks/IconDetector.java

@@ -1985,7 +1985,7 @@ public class IconDetector extends ResourceXmlDetector implements UastScanner {
         public boolean visitFunction(@NotNull UFunction node) {
             if (node.matchesName(ON_CREATE_OPTIONS_MENU)) {
                 // Gather any R.menu references found in this method
-                new MenuFinder().process(node);
+                node.accept(new MenuFinder());
             }
 
             return super.visitFunction(node);
@@ -2036,7 +2036,7 @@ public class IconDetector extends ResourceXmlDetector implements UastScanner {
                     UFunction method = UastUtils.getContainingFunction(node);
                     if (method != null) {
                         SetIconFinder finder = new SetIconFinder();
-                        finder.process(method);
+                        method.accept(finder);
                     }
                 }
             }

plugins/lint/lint-checks/src/com/android/tools/klint/checks/JavaPerformanceDetector.java

@@ -324,7 +324,7 @@ public class JavaPerformanceDetector extends Detector implements UastScanner {
                     UExpression thenBranch = ifNode.getThenBranch();
                     if (thenBranch != null) {
                         AssignmentTracker visitor = new AssignmentTracker(assignments);
-                        visitor.process(thenBranch);
+                        thenBranch.accept(visitor);
                     }
                     if (!assignments.isEmpty()) {
                         List<String> references = new ArrayList<String>();

plugins/lint/lint-checks/src/com/android/tools/klint/checks/JavaScriptInterfaceDetector.java

@@ -109,7 +109,7 @@ public class JavaScriptInterfaceDetector extends Detector implements UastScanner
             UFunction method = UastUtils.getContainingFunction(node);
             if (method != null) {
                 ConcreteTypeVisitor v = new ConcreteTypeVisitor(context, node);
-                v.process(method);
+                method.accept(v);
                 resolved = v.getType();
                 if (resolved == null) {
                     return;
@@ -211,10 +211,8 @@ public class JavaScriptInterfaceDetector extends Detector implements UastScanner
         }
 
         @Override
-        public void process(@NotNull UElement element) {
+        public boolean visitElement(@NotNull UElement node) {
-            if (!mFoundCall) {
+            return mFoundCall || super.visitElement(node);
-                super.process(element);
-            }
         }
 
         @Override

plugins/lint/lint-checks/src/com/android/tools/klint/checks/PermissionRequirement.java

@@ -590,13 +590,13 @@ public abstract class PermissionRequirement {
 
             if (node != null) {
                 final AtomicReference<UExpression> reference = new AtomicReference<UExpression>();
-                new UastVisitor() {
+                node.accept(new UastVisitor() {
                     @Override
                     public boolean visitVariable(@NotNull UVariable node) {
                         reference.set(node.getInitializer());
                         return true;
                     }
-                }.process(node);
+                });
                 UExpression expression = reference.get();
                 if (expression != null) {
                     return parse(annotation, expression);

plugins/lint/lint-checks/src/com/android/tools/klint/checks/SecurityDetector.java

@@ -389,10 +389,10 @@ public class SecurityDetector extends Detector implements Detector.XmlScanner, U
         public boolean visitCallExpression(@NotNull UCallExpression node) {
             String name = node.getFunctionName();
             if ("openFileOutput".equals(name) || "getSharedPreferences".equals(name)) {
-                identifierHandler.processChildren(node);
+                return false;
             }
 
-            return false;
+            return true;
         }
     }
 }

plugins/lint/lint-checks/src/com/android/tools/klint/checks/SharedPrefsDetector.java

@@ -147,7 +147,7 @@ public class SharedPrefsDetector extends Detector implements UastScanner {
         }
 
         CommitFinder finder = new CommitFinder(context, node, allowCommitBeforeTarget);
-        finder.process(method);
+        method.accept(finder);
         if (!finder.isCommitCalled()) {
             context.report(ISSUE, method, UastAndroidUtils.getLocation(node),
                            "`SharedPreferences.edit()` without a corresponding `commit()` or `apply()` call");

plugins/lint/lint-checks/src/com/android/tools/klint/checks/StringFormatDetector.java

@@ -1043,7 +1043,7 @@ public class StringFormatDetector extends ResourceXmlDetector implements UastSca
 
         JavaContext lintContext = context.getLintContext();
         UastStringTracker tracker = new UastStringTracker(lintContext, method, call, 0);
-        tracker.process(method);
+        method.accept(tracker);
         String name = tracker.getFormatStringName();
         if (name == null) {
             return;
@@ -1314,7 +1314,7 @@ public class StringFormatDetector extends ResourceXmlDetector implements UastSca
       @NonNull UElement call) {
         assert call instanceof UCallExpression;
         UastStringTracker tracker = new UastStringTracker(null, method, call, 0);
-        tracker.process(method);
+        method.accept(tracker);
 
         return tracker.getFormatStringName();
     }
@@ -1327,7 +1327,7 @@ public class StringFormatDetector extends ResourceXmlDetector implements UastSca
       int argIndex) {
         assert call instanceof UCallExpression;
         UastStringTracker tracker = new UastStringTracker(null, method, call, argIndex);
-        tracker.process(method);
+        method.accept(tracker);
 
         return tracker.getFormatStringName();
     }
@@ -1481,10 +1481,8 @@ public class StringFormatDetector extends ResourceXmlDetector implements UastSca
         }
 
         @Override
-        public void process(@NotNull UElement element) {
+        public boolean visitElement(@NotNull UElement node) {
-            if (!mDone) {
+            return mDone || super.visitElement(node);
-                super.process(element);
-            }
         }
 
         @Override

plugins/lint/lint-checks/src/com/android/tools/klint/checks/SupportAnnotationDetector.java

@@ -388,7 +388,7 @@ public class SupportAnnotationDetector extends Detector implements UastScanner {
                 UFunction methodNode = UastUtils.getContainingFunction(node);
                 if (methodNode != null) {
                     CheckPermissionVisitor visitor = new CheckPermissionVisitor(node);
-                    visitor.process(methodNode);
+                    methodNode.accept(visitor);
                     handlesMissingPermission = visitor.checksPermission();
                 }
             }
@@ -466,10 +466,8 @@ public class SupportAnnotationDetector extends Detector implements UastScanner {
         }
 
         @Override
-        public void process(@NotNull UElement element) {
+        public boolean visitElement(@NotNull UElement node) {
-            if (!mDone) {
+            return mDone;
-                super.process(element);
-            }
         }
 
         @Override

plugins/lint/lint-checks/src/com/android/tools/klint/checks/ToastDetector.java

@@ -105,7 +105,7 @@ public class ToastDetector extends Detector implements UastScanner {
 
         UExpression nodeWithPossibleQualifier = UastUtils.getQualifiedCallElement(node);
         ShowFinder finder = new ShowFinder(nodeWithPossibleQualifier);
-        finder.process(method);
+        method.accept(finder);
         if (!finder.isShowCalled()) {
             context.report(ISSUE, node, UastAndroidUtils.getLocation(node),
                            "Toast created but not shown: did you forget to call `show()` ?");

plugins/lint/lint-checks/src/com/android/tools/klint/checks/ViewHolderDetector.java

@@ -95,7 +95,7 @@ public class ViewHolderDetector extends Detector implements UastScanner {
         public boolean visitFunction(@NotNull UFunction node) {
             if (isViewAdapterMethod(node)) {
                 InflationVisitor visitor = new InflationVisitor(mContext);
-                visitor.process(node);
+                node.accept(visitor);
                 visitor.finish();
             }
             return false;

plugins/lint/uast-android/src/org/jetbrains/uast/check/UastChecker.kt

@@ -38,26 +38,10 @@ interface UastAndroidContext : UastContext {
 }
 
 object UastChecker {
-    fun checkWithCustomHandler(
+    fun check(project: Project, file: File, context: UastAndroidContext, visitor: UastVisitor) {
-            project: Project,
-            file: File,
-            context: UastAndroidContext,
-            visitor: UastVisitor) {
-        check(project, file, context, UastCallback { visitor.handle(it) })
-    }
-
-    fun check(project: Project, file: File, context: UastAndroidContext, callback: UastCallback) {
         val vfile = VirtualFileManager.getInstance().findFileByUrl("file://" + file.absolutePath) ?: return
 
         val plugins = context.languagePlugins
-        val additionalCheckers = plugins.fold(mutableListOf<UastAdditionalChecker>()) { list, plugin ->
-            for (checker in plugin.additionalCheckers) {
-                if (checker is AndroidUastAdditionalChecker) list += checker
-            }
-            list
-        }
-
-        val handlerWrapper = CallbackWrapper(callback, additionalCheckers, context)
 
         ApplicationManager.getApplication().runReadAction {
             val psiFile = PsiManager.getInstance(project).findFile(vfile)
@@ -66,12 +50,12 @@ object UastChecker {
                 when (psiFile) {
                     is PsiJavaFile -> {
                         val ufile = JavaUastLanguagePlugin.converter.convertWithParent(psiFile)
-                        ufile?.handleTraverse(handlerWrapper)
+                        ufile?.accept(visitor)
                     }
                     else -> for (plugin in plugins) {
                         val ufile = plugin.converter.convertWithParent(psiFile)
                         if (ufile != null) {
-                            ufile.handleTraverse(handlerWrapper)
+                            ufile.accept(visitor)
                             break
                         }
                     }
@@ -80,19 +64,6 @@ object UastChecker {
         }
     }
 
-    private class CallbackWrapper(
-            val original: UastCallback,
-            val additionalCheckers: List<UastAdditionalChecker>,
-            val context: UastAndroidContext
-    ) : UastCallback {
-        override fun invoke(element: UElement) {
-            original(element)
-            for (checker in additionalCheckers) {
-                checker(element, this, context)
-            }
-        }
-    }
-
     fun check(project: Project, file: File, scanner: UastScanner, context: UastAndroidContext) {
         val applicableFunctionNames = scanner.applicableFunctionNames ?: emptyList()
         val applicableSuperClasses = scanner.applicableSuperClasses ?: emptyList()
@@ -100,57 +71,64 @@ object UastChecker {
 
         val appliesToResourcesRefs = scanner.appliesToResourceRefs()
 
-        var callback: UastCallback?
+        val visitor = object : UastVisitor() {
-        callback = UastCallback { element ->
+            override fun visitCallExpression(node: UCallExpression): Boolean {
-            when (element) {
+                if (applicableFunctionNames.isNotEmpty()) {
-                is UCallExpression -> {
+                    if (node.kind == FUNCTION_CALL && node.functionName in applicableFunctionNames) {
-                    if (applicableFunctionNames.isNotEmpty()) {
+                        scanner.visitFunctionCall(context, node)
-                        if (element.kind == FUNCTION_CALL && element.functionName in applicableFunctionNames) {
-                            scanner.visitFunctionCall(context, element)
-                        }
                     }
+                }
 
-                    if (applicableConstructorTypes.isNotEmpty()) {
+                if (applicableConstructorTypes.isNotEmpty()) {
-                        if (element.kind == CONSTRUCTOR_CALL) {
+                    if (node.kind == CONSTRUCTOR_CALL) {
-                            element.resolve(context)?.let { constructor ->
+                        node.resolve(context)?.let { constructor ->
-                                if (constructor.getContainingClass()?.fqName in applicableConstructorTypes) {
+                            if (constructor.getContainingClass()?.fqName in applicableConstructorTypes) {
-                                    scanner.visitConstructor(context, element, constructor)
+                                scanner.visitConstructor(context, node, constructor)
-                                }
                             }
                         }
                     }
                 }
-                is UClass -> if (applicableSuperClasses.isNotEmpty()) {
+
-                    if (applicableSuperClasses.any { element.isSubclassOf(it) }) {
+                return false
-                        scanner.visitClass(context, element)
+            }
+
+            override fun visitClass(node: UClass): Boolean {
+                if (applicableSuperClasses.isNotEmpty()) {
+                    if (applicableSuperClasses.any { node.isSubclassOf(it) }) {
+                        scanner.visitClass(context, node)
                     }
                 }
-                is UQualifiedExpression -> {
-                    if (appliesToResourcesRefs && element.receiver is UQualifiedExpression) {
-                        val parentQualifiedExpr = element.receiver as UQualifiedExpression
-                        val resourceName = element.selector
-                        val resourceType = parentQualifiedExpr.selector
-                        val receiver = parentQualifiedExpr.receiver
 
-                        val receiverIsResourceClass = when (receiver) {
+                return false
-                            is USimpleReferenceExpression -> receiver.identifier == "R"
+            }
-                            is UQualifiedExpression -> receiver.selectorMatches("R")
-                            else -> false
-                        }
 
-                        if (resourceName is USimpleReferenceExpression && resourceType is USimpleReferenceExpression
+            override fun visitQualifiedExpression(node: UQualifiedExpression): Boolean {
-                                && receiverIsResourceClass && receiver is UResolvable) {
+                if (appliesToResourcesRefs && node.receiver is UQualifiedExpression) {
-                            val resolvedReceiver = receiver.resolve(context)
+                    val parentQualifiedExpr = node.receiver as UQualifiedExpression
-                            val isFramework = (resolvedReceiver as? UClass)?.matchesFqName("android.R") ?: false
+                    val resourceName = node.selector
+                    val resourceType = parentQualifiedExpr.selector
+                    val receiver = parentQualifiedExpr.receiver
 
-                            scanner.visitResourceReference(context, element, resourceType.identifier, resourceName.identifier, isFramework)
+                    val receiverIsResourceClass = when (receiver) {
-                        }
+                        is USimpleReferenceExpression -> receiver.identifier == "R"
+                        is UQualifiedExpression -> receiver.selectorMatches("R")
+                        else -> false
+                    }
+
+                    if (resourceName is USimpleReferenceExpression && resourceType is USimpleReferenceExpression
+                        && receiverIsResourceClass && receiver is UResolvable) {
+                        val resolvedReceiver = receiver.resolve(context)
+                        val isFramework = (resolvedReceiver as? UClass)?.matchesFqName("android.R") ?: false
+
+                        scanner.visitResourceReference(context, node, resourceType.identifier, resourceName.identifier, isFramework)
                     }
                 }
+
+                return false
             }
         }
 
-        check(project, file, context, callback)
+        check(project, file, context, visitor)
     }
 
 }

plugins/uast-common/src/org/jetbrains/uast/UastUtils.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.uast
 
 import org.jetbrains.uast.kinds.UastClassKind
+import org.jetbrains.uast.visitor.UastVisitor
 
 tailrec fun UElement?.getContainingClass(): UClass? {
     val parent = this?.parent ?: return null
@@ -44,11 +45,6 @@ tailrec fun UElement?.getContainingDeclaration(): UDeclaration? {
     return parent.getContainingDeclaration()
 }
 
-fun UElement.handleTraverse(callback: UastCallback) {
-    callback(this)
-    this.traverse(callback)
-}
-
 fun UDeclaration.isTopLevel() = parent is UFile
 
 fun UElement.log(firstLine: String, vararg nested: Any?): String {
@@ -67,10 +63,9 @@ fun UElement.log(firstLine: String, vararg nested: Any?): String {
 val String.withMargin: String
     get() = lines().joinToString("\n") { "    " + it }
 
-fun List<UElement>.handleTraverseList(callback: UastCallback) {
+fun List<UElement>.acceptList(visitor: UastVisitor) {
     for (element in this) {
-        callback(element)
+        element.accept(visitor)
-        element.traverse(callback)
     }
 }
 

plugins/uast-common/src/org/jetbrains/uast/UastVisitor.kt

@@ -18,108 +18,51 @@ package org.jetbrains.uast.visitor
 import org.jetbrains.uast.*
 
 abstract class UastVisitor {
-    open fun visitFile(node: UFile): Boolean = false
+    open fun visitElement(node: UElement): Boolean = false
+
+    open fun visitFile(node: UFile) = visitElement(node)
+    open fun visitImportStatement(node: UImportStatement) = visitElement(node)
+    open fun visitAnnotation(node: UAnnotation) = visitElement(node)
+    open fun visitCatchClause(node: UCatchClause) = visitElement(node)
+    open fun visitType(node: UType) = visitElement(node)
 
     // Declarations
-    open fun visitClass(node: UClass): Boolean = false
+    open fun visitClass(node: UClass) = visitElement(node)
-    open fun visitFunction(node: UFunction): Boolean = false
+    open fun visitFunction(node: UFunction) = visitElement(node)
-    open fun visitVariable(node: UVariable): Boolean = false
+    open fun visitVariable(node: UVariable) = visitElement(node)
-    open fun visitImportStatement(node: UImportStatement): Boolean = false
+
-    open fun visitAnnotation(node: UAnnotation): Boolean = false
 
     // Expressions
-    open fun visitLabeledExpression(node: ULabeledExpression): Boolean = false
+    open fun visitLabeledExpression(node: ULabeledExpression) = visitElement(node)
-    open fun visitDeclarationsExpression(node: UDeclarationsExpression): Boolean = false
+    open fun visitDeclarationsExpression(node: UDeclarationsExpression) = visitElement(node)
-    open fun visitBlockExpression(node: UBlockExpression): Boolean = false
+    open fun visitBlockExpression(node: UBlockExpression) = visitElement(node)
-    open fun visitQualifiedExpression(node: UQualifiedExpression): Boolean = false
+    open fun visitQualifiedExpression(node: UQualifiedExpression) = visitElement(node)
-    open fun visitSimpleReferenceExpression(node: USimpleReferenceExpression): Boolean = false
+    open fun visitSimpleReferenceExpression(node: USimpleReferenceExpression) = visitElement(node)
-    open fun visitCallExpression(node: UCallExpression): Boolean = false
+    open fun visitCallExpression(node: UCallExpression) = visitElement(node)
-    open fun visitBinaryExpression(node: UBinaryExpression): Boolean = false
+    open fun visitBinaryExpression(node: UBinaryExpression) = visitElement(node)
-    open fun visitBinaryExpressionWithType(node: UBinaryExpressionWithType): Boolean = false
+    open fun visitBinaryExpressionWithType(node: UBinaryExpressionWithType) = visitElement(node)
-    open fun visitParenthesizedExpression(node: UParenthesizedExpression): Boolean = false
+    open fun visitParenthesizedExpression(node: UParenthesizedExpression) = visitElement(node)
-    open fun visitPrefixExpression(node: UPrefixExpression): Boolean = false
+    open fun visitUnaryExpression(node: UUnaryExpression) = visitElement(node)
-    open fun visitPostfixExpression(node: UPostfixExpression): Boolean = false
+    open fun visitPrefixExpression(node: UPrefixExpression) = visitElement(node)
-    open fun visitSpecialExpressionList(node: USpecialExpressionList): Boolean = false
+    open fun visitPostfixExpression(node: UPostfixExpression) = visitElement(node)
-    open fun visitExpressionList(node: UExpressionList): Boolean = false
+    open fun visitSpecialExpressionList(node: USpecialExpressionList) = visitElement(node)
-    open fun visitIfExpression(node: UIfExpression): Boolean = false
+    open fun visitExpressionList(node: UExpressionList) = visitElement(node)
-    open fun visitSwitchExpression(node: USwitchExpression): Boolean = false
+    open fun visitIfExpression(node: UIfExpression) = visitElement(node)
-    open fun visitSwitchClauseExpression(node: USwitchClauseExpression): Boolean = false
+    open fun visitSwitchExpression(node: USwitchExpression) = visitElement(node)
-    open fun visitWhileExpression(node: UWhileExpression): Boolean = false
+    open fun visitSwitchClauseExpression(node: USwitchClauseExpression) = visitElement(node)
-    open fun visitDoWhileExpression(node: UDoWhileExpression): Boolean = false
+    open fun visitWhileExpression(node: UWhileExpression) = visitElement(node)
-    open fun visitForExpression(node: UForExpression): Boolean = false
+    open fun visitDoWhileExpression(node: UDoWhileExpression) = visitElement(node)
-    open fun visitForEachExpression(node: UForEachExpression): Boolean = false
+    open fun visitForExpression(node: UForExpression) = visitElement(node)
-    open fun visitTryExpression(node: UTryExpression): Boolean = false
+    open fun visitForEachExpression(node: UForEachExpression) = visitElement(node)
-    open fun visitLiteralExpression(node: ULiteralExpression): Boolean = false
+    open fun visitTryExpression(node: UTryExpression) = visitElement(node)
-    open fun visitThisExpression(node: UThisExpression): Boolean = false
+    open fun visitLiteralExpression(node: ULiteralExpression) = visitElement(node)
-    open fun visitSuperExpression(node: USuperExpression): Boolean = false
+    open fun visitThisExpression(node: UThisExpression) = visitElement(node)
-    open fun visitArrayAccessExpression(node: UArrayAccessExpression): Boolean = false
+    open fun visitSuperExpression(node: USuperExpression) = visitElement(node)
-    open fun visitCallableReferenceExpression(node: UCallableReferenceExpression): Boolean = false
+    open fun visitArrayAccessExpression(node: UArrayAccessExpression) = visitElement(node)
-    open fun visitClassLiteralExpression(node: UClassLiteralExpression): Boolean = false
+    open fun visitCallableReferenceExpression(node: UCallableReferenceExpression) = visitElement(node)
-    open fun visitLambdaExpression(node: ULambdaExpression): Boolean = false
+    open fun visitClassLiteralExpression(node: UClassLiteralExpression) = visitElement(node)
-    open fun visitObjectLiteralExpression(node: UObjectLiteralExpression): Boolean = false
+    open fun visitLambdaExpression(node: ULambdaExpression) = visitElement(node)
-
+    open fun visitObjectLiteralExpression(node: UObjectLiteralExpression) = visitElement(node)
-    open fun beforeVisit(node: UElement) {}
-    open fun visitOther(node: UElement): Boolean = true
-    open fun afterVisit(node: UElement) {}
-
-    fun handle(node: UElement): Boolean {
-        beforeVisit(node)
-        val result = when (node) {
-            is UFile -> visitFile(node)
-
-            is UClass -> visitClass(node)
-            is UFunction -> visitFunction(node)
-            is UVariable -> visitVariable(node)
-            is UImportStatement -> visitImportStatement(node)
-            is UAnnotation -> visitAnnotation(node)
-
-            is ULabeledExpression -> visitLabeledExpression(node)
-            is UDeclarationsExpression -> visitDeclarationsExpression(node)
-            is UBlockExpression -> visitBlockExpression(node)
-            is UQualifiedExpression -> visitQualifiedExpression(node)
-            is USimpleReferenceExpression -> visitSimpleReferenceExpression(node)
-            is UCallExpression -> visitCallExpression(node)
-            is UBinaryExpression -> visitBinaryExpression(node)
-            is UBinaryExpressionWithType -> visitBinaryExpressionWithType(node)
-            is UParenthesizedExpression -> visitParenthesizedExpression(node)
-            is UPrefixExpression -> visitPrefixExpression(node)
-            is UPostfixExpression -> visitPostfixExpression(node)
-            is USpecialExpressionList -> visitSpecialExpressionList(node)
-            is UExpressionList -> visitExpressionList(node)
-            is UIfExpression -> visitIfExpression(node)
-            is USwitchExpression -> visitSwitchExpression(node)
-            is USwitchClauseExpression -> visitSwitchClauseExpression(node)
-            is UWhileExpression -> visitWhileExpression(node)
-            is UDoWhileExpression -> visitDoWhileExpression(node)
-            is UForExpression -> visitForExpression(node)
-            is UForEachExpression -> visitForEachExpression(node)
-            is UTryExpression -> visitTryExpression(node)
-            is ULiteralExpression -> visitLiteralExpression(node)
-            is UThisExpression -> visitThisExpression(node)
-            is USuperExpression -> visitSuperExpression(node)
-            is UArrayAccessExpression -> visitArrayAccessExpression(node)
-            is UCallableReferenceExpression -> visitCallableReferenceExpression(node)
-            is UClassLiteralExpression -> visitClassLiteralExpression(node)
-            is ULambdaExpression -> visitLambdaExpression(node)
-            is UObjectLiteralExpression -> visitObjectLiteralExpression(node)
-
-            else -> visitOther(node)
-        }
-        afterVisit(node)
-        return result
-    }
-    
-    open fun process(element: UElement) {
-        if (!handle(element)) {
-            processChildren(element)
-        }
-    }
-
-    fun processChildren(element: UElement) {
-        element.traverse(UastCallback { handle(it) })
-    }
 }
 
-object EmptyUastVisitor : UastVisitor() {
+object EmptyUastVisitor : UastVisitor()
-    override fun process(element: UElement) {}
-}

plugins/uast-common/src/org/jetbrains/uast/baseElements/UAnnotation.kt

@@ -15,10 +15,12 @@
  */
 package org.jetbrains.uast
 
-interface UAnnotation : UElement, UNamed, UFqNamed, LeafUElement {
+import org.jetbrains.uast.visitor.UastVisitor
-    fun resolve(context: UastContext): UClass?
 
+interface UAnnotation : UElement, UNamed, UFqNamed {
     val valueArguments: List<UNamedExpression>
+
+    fun resolve(context: UastContext): UClass?
     fun getValue(name: String) = valueArguments.firstOrNull { it.name == name }?.expression?.evaluate()
 
     fun getValues() = valueArguments.map { Pair(it.name, it.expression.evaluate()) }
@@ -26,6 +28,11 @@ interface UAnnotation : UElement, UNamed, UFqNamed, LeafUElement {
     override fun logString() = log("UAnnotation ($name)")
     override fun renderString() = if (valueArguments.isEmpty()) "@$name" else "@$name(" +
             valueArguments.joinToString { it.renderString() } + ")"
+
+    override fun accept(visitor: UastVisitor) {
+        if (visitor.visitAnnotation(this)) return
+        valueArguments.acceptList(visitor)
+    }
 }
 
 interface UAnnotated : UElement {

plugins/uast-common/src/org/jetbrains/uast/baseElements/UElement.kt

@@ -15,6 +15,8 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UElement {
     val parent: UElement?
 
@@ -23,7 +25,10 @@ interface UElement {
     
     fun logString(): String
     fun renderString(): String = logString()
-    fun traverse(callback: UastCallback)
+
+    fun accept(visitor: UastVisitor) {
+        visitor.visitElement(this)
+    }
 }
 
 interface UNamed {
@@ -42,10 +47,6 @@ interface UModifierOwner {
     fun hasModifier(modifier: UastModifier): Boolean
 }
 
-interface LeafUElement : UElement {
-    override fun traverse(callback: UastCallback) {}
-}
-
 interface UResolvable {
     fun resolve(context: UastContext): UDeclaration?
     fun resolveOrEmpty(context: UastContext): UDeclaration = resolve(context) ?: UDeclarationNotResolved

plugins/uast-common/src/org/jetbrains/uast/baseElements/UExpression.kt

@@ -15,10 +15,16 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UExpression : UElement {
     open fun evaluate(): Any? = null
     fun evaluateString(): String? = evaluate() as? String
     fun getExpressionType(): UType? = null
+
+    override fun accept(visitor: UastVisitor) {
+        visitor.visitElement(this)
+    }
 }
 
 interface NoAnnotations : UAnnotated {
@@ -30,6 +36,6 @@ interface NoModifiers : UModifierOwner {
     override fun hasModifier(modifier: UastModifier) = false
 }
 
-class EmptyExpression(override val parent: UElement) : UExpression, LeafUElement {
+class EmptyExpression(override val parent: UElement) : UExpression {
     override fun logString() = "EmptyExpression"
 }

plugins/uast-common/src/org/jetbrains/uast/baseElements/UFile.kt

@@ -15,6 +15,8 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UFile: UElement {
     val packageFqName: String?
     val importStatements: List<UImportStatement>
@@ -26,9 +28,10 @@ interface UFile: UElement {
     override val parent: UElement?
         get() = null
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        declarations.handleTraverseList(callback)
+        if (visitor.visitFile(this)) return
-        importStatements.handleTraverseList(callback)
+        declarations.acceptList(visitor)
+        importStatements.acceptList(visitor)
     }
 
     override fun logString() = "UFile (package = $packageFqName)\n" + declarations.joinToString("\n") { it.logString().withMargin }

plugins/uast-common/src/org/jetbrains/uast/controlStructures/UDoWhileExpression.kt

@@ -15,12 +15,15 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UDoWhileExpression : ULoopExpression {
     val condition: UExpression
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        condition.handleTraverse(callback)
+        if (visitor.visitDoWhileExpression(this)) return
-        body.handleTraverse(callback)
+        condition.accept(visitor)
+        body.accept(visitor)
     }
 
     override fun renderString() = buildString {

plugins/uast-common/src/org/jetbrains/uast/controlStructures/UForEachExpression.kt

@@ -15,14 +15,17 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UForEachExpression : ULoopExpression {
     val variable: UVariable
     val iteratedValue: UExpression
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        variable.handleTraverse(callback)
+        if (visitor.visitForEachExpression(this)) return
-        iteratedValue.handleTraverse(callback)
+        variable.accept(visitor)
-        body.handleTraverse(callback)
+        iteratedValue.accept(visitor)
+        body.accept(visitor)
     }
 
     override fun renderString() = buildString {

plugins/uast-common/src/org/jetbrains/uast/controlStructures/UForExpression.kt

@@ -15,16 +15,19 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UForExpression : ULoopExpression {
     val declaration: UExpression?
     val condition: UExpression?
     val update: UExpression?
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        declaration?.handleTraverse(callback)
+        if (visitor.visitForExpression(this)) return
-        condition?.handleTraverse(callback)
+        declaration?.accept(visitor)
-        update?.handleTraverse(callback)
+        condition?.accept(visitor)
-        body.handleTraverse(callback)
+        update?.accept(visitor)
+        body.accept(visitor)
     }
 
     override fun renderString() = buildString {

plugins/uast-common/src/org/jetbrains/uast/controlStructures/UIfExpression.kt

@@ -15,16 +15,19 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UIfExpression : UExpression {
     val condition: UExpression
     val thenBranch: UExpression?
     val elseBranch: UExpression?
     val isTernary: Boolean
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        condition.handleTraverse(callback)
+        if (visitor.visitIfExpression(this)) return
-        thenBranch?.handleTraverse(callback)
+        condition.accept(visitor)
-        elseBranch?.handleTraverse(callback)
+        thenBranch?.accept(visitor)
+        elseBranch?.accept(visitor)
     }
 
     override fun logString() = log("UIfExpression", condition, thenBranch, elseBranch)

plugins/uast-common/src/org/jetbrains/uast/controlStructures/USwitchExpression.kt

@@ -15,13 +15,16 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface USwitchExpression : UExpression {
     val expression: UExpression?
     val body: UExpression
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        expression?.handleTraverse(callback)
+        if (visitor.visitSwitchExpression(this)) return
-        body.handleTraverse(callback)
+        expression?.accept(visitor)
+        body.accept(visitor)
     }
 
     override fun logString() = log("USwitchExpression", expression, body)
@@ -37,8 +40,9 @@ interface USwitchClauseExpression : UExpression
 interface UExpressionSwitchClauseExpression : USwitchClauseExpression {
     val caseValue: UExpression
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        caseValue.handleTraverse(callback)
+        if (visitor.visitSwitchClauseExpression(this)) return
+        caseValue.accept(visitor)
     }
 
     override fun renderString() = caseValue.renderString() + " -> "
@@ -46,7 +50,6 @@ interface UExpressionSwitchClauseExpression : USwitchClauseExpression {
 }
 
 interface UDefaultSwitchClauseExpression : USwitchClauseExpression {
-    override fun traverse(callback: UastCallback) {}
     override fun logString() = "UDefaultSwitchClause"
     override fun renderString() = "else -> "
 }

plugins/uast-common/src/org/jetbrains/uast/controlStructures/UTryExpression.kt

@@ -15,17 +15,20 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UTryExpression : UExpression {
     val resources: List<UElement>?
     val tryClause: UExpression
     val catchClauses: List<UCatchClause>
     val finallyClause: UExpression?
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        resources?.handleTraverseList(callback)
+        if (visitor.visitTryExpression(this)) return
-        tryClause.handleTraverse(callback)
+        resources?.acceptList(visitor)
-        catchClauses.handleTraverseList(callback)
+        tryClause.accept(visitor)
-        finallyClause?.handleTraverse(callback)
+        catchClauses.acceptList(visitor)
+        finallyClause?.accept(visitor)
     }
 
     override fun renderString() = buildString {
@@ -46,10 +49,11 @@ interface UCatchClause : UElement {
     val parameters: List<UVariable>
     val types: List<UType>
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        body.handleTraverse(callback)
+        if (visitor.visitCatchClause(this)) return
-        parameters.handleTraverseList(callback)
+        body.accept(visitor)
-        types.handleTraverseList(callback)
+        parameters.acceptList(visitor)
+        types.acceptList(visitor)
     }
 
     override fun logString() = log("UCatchClause", body)

plugins/uast-common/src/org/jetbrains/uast/controlStructures/UWhileExpression.kt

@@ -15,12 +15,15 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UWhileExpression : ULoopExpression {
     val condition: UExpression
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        condition.handleTraverse(callback)
+        if (visitor.visitWhileExpression(this)) return
-        body.handleTraverse(callback)
+        condition.accept(visitor)
+        body.accept(visitor)
     }
 
     override fun renderString() = buildString {

plugins/uast-common/src/org/jetbrains/uast/declarations/UClass.kt

@@ -16,6 +16,7 @@
 package org.jetbrains.uast
 
 import org.jetbrains.uast.kinds.UastClassKind
+import org.jetbrains.uast.visitor.UastVisitor
 
 interface UClass : UDeclaration, UFqNamed, UModifierOwner, UAnnotated {
     /* The simple class name is only for the debug purposes. Do not check against it in the production code */
@@ -51,10 +52,11 @@ interface UClass : UDeclaration, UFqNamed, UModifierOwner, UAnnotated {
 
     fun getSuperClass(context: UastContext): UClass?
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        nameElement?.handleTraverse(callback)
+        if (visitor.visitClass(this)) return
-        declarations.handleTraverseList(callback)
+        nameElement?.accept(visitor)
-        annotations.handleTraverseList(callback)
+        declarations.acceptList(visitor)
+        annotations.acceptList(visitor)
     }
 
     override fun renderString(): String {

plugins/uast-common/src/org/jetbrains/uast/declarations/UDeclaration.kt

@@ -15,11 +15,17 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UDeclaration : UElement, UNamed {
     val nameElement: UElement?
+
+    override fun accept(visitor: UastVisitor) {
+        visitor.visitElement(this)
+    }
 }
 
-object UDeclarationNotResolved : UDeclaration, LeafUElement {
+object UDeclarationNotResolved : UDeclaration {
     override val name = "<declaration not resolved>"
     override val nameElement = null
     override val parent = null

plugins/uast-common/src/org/jetbrains/uast/declarations/UFunction.kt

@@ -15,6 +15,8 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UFunction : UDeclaration, UModifierOwner, UAnnotated {
     val kind: UastFunctionKind
     val valueParameters: List<UVariable>
@@ -30,13 +32,14 @@ interface UFunction : UDeclaration, UModifierOwner, UAnnotated {
 
     fun getSuperFunctions(context: UastContext): List<UFunction>
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        nameElement?.handleTraverse(callback)
+        if (visitor.visitFunction(this)) return
-        valueParameters.handleTraverseList(callback)
+        nameElement?.accept(visitor)
-        body.handleTraverse(callback)
+        valueParameters.acceptList(visitor)
-        annotations.handleTraverseList(callback)
+        body.accept(visitor)
-        typeParameters.handleTraverseList(callback)
+        annotations.acceptList(visitor)
-        returnType?.handleTraverse(callback)
+        typeParameters.acceptList(visitor)
+        returnType?.accept(visitor)
     }
 
     override fun renderString(): String {

plugins/uast-common/src/org/jetbrains/uast/declarations/UImportStatement.kt

@@ -16,6 +16,7 @@
 package org.jetbrains.uast
 
 import org.jetbrains.uast.kinds.UastImportKind
+import org.jetbrains.uast.visitor.UastVisitor
 
 interface UImportStatement : UElement {
     val fqNameToImport: String?
@@ -25,7 +26,10 @@ interface UImportStatement : UElement {
     /* Returns null if isStarImport = true */
     fun resolve(context: UastContext): UDeclaration?
 
-    override fun traverse(callback: UastCallback) {}
+    override fun accept(visitor: UastVisitor) {
+        visitor.visitImportStatement(this)
+    }
+
     override fun logString() = "UImport ($fqNameToImport)"
     override fun renderString() = "import ${fqNameToImport ?: ""}"
 }

plugins/uast-common/src/org/jetbrains/uast/declarations/UType.kt

@@ -15,7 +15,9 @@
  */
 package org.jetbrains.uast
 
-interface UType : UElement, UNamed, UFqNamed, UAnnotated, UResolvable, LeafUElement {
+import org.jetbrains.uast.visitor.UastVisitor
+
+interface UType : UElement, UNamed, UFqNamed, UAnnotated, UResolvable {
     override fun matchesName(name: String) = this.name == name || this.name.endsWith(".$name")
 
     /* The simple type name is only for the debug purposes. Do not check against it in the production code */
@@ -36,12 +38,13 @@ interface UType : UElement, UNamed, UFqNamed, UAnnotated, UResolvable, LeafUElem
     override fun resolve(context: UastContext): UClass?
     override fun resolveOrEmpty(context: UastContext) = resolve(context) ?: UClassNotResolved
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        annotations.handleTraverseList(callback)
+        if (visitor.visitType(this)) return
+        annotations.acceptList(visitor)
     }
 }
 
-interface UTypeReference : UDeclaration, UResolvable, LeafUElement {
+interface UTypeReference : UDeclaration, UResolvable {
     override fun renderString() = ""
     override fun logString() = log("UTypeReference")
 

plugins/uast-common/src/org/jetbrains/uast/declarations/UVariable.kt

@@ -15,6 +15,8 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UVariable : UDeclaration, UModifierOwner, UAnnotated {
     val initializer: UExpression?
     val kind: UastVariableKind
@@ -27,11 +29,12 @@ interface UVariable : UDeclaration, UModifierOwner, UAnnotated {
     open val setters: List<UFunction>?
         get() = null
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        nameElement?.handleTraverse(callback)
+        if (visitor.visitVariable(this)) return
-        initializer?.handleTraverse(callback)
+        nameElement?.accept(visitor)
-        annotations.handleTraverseList(callback)
+        initializer?.accept(visitor)
-        type.handleTraverse(callback)
+        annotations.acceptList(visitor)
+        type.accept(visitor)
     }
 
     override fun renderString(): String {

plugins/uast-common/src/org/jetbrains/uast/expressions/UArrayAccessExpression.kt

@@ -15,13 +15,16 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UArrayAccessExpression : UExpression {
     val receiver: UExpression
     val indices: List<UExpression>
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        receiver.handleTraverse(callback)
+        if (visitor.visitArrayAccessExpression(this)) return
-        indices.handleTraverseList(callback)
+        receiver.accept(visitor)
+        indices.acceptList(visitor)
     }
 
     override fun logString() = log("UArrayAccessExpression", receiver, indices)

plugins/uast-common/src/org/jetbrains/uast/expressions/UBinaryExpression.kt

@@ -15,14 +15,17 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UBinaryExpression : UExpression {
     val leftOperand: UExpression
     val operator: UastBinaryOperator
     val rightOperand: UExpression
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        leftOperand.handleTraverse(callback)
+        if (visitor.visitBinaryExpression(this)) return
-        rightOperand.handleTraverse(callback)
+        leftOperand.accept(visitor)
+        rightOperand.accept(visitor)
     }
 
     override fun logString() =

plugins/uast-common/src/org/jetbrains/uast/expressions/UBinaryExpressionWithType.kt

@@ -15,6 +15,8 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UBinaryExpressionWithType : UExpression {
     val operand: UExpression
     val operationKind: UastBinaryExpressionWithTypeKind
@@ -23,8 +25,9 @@ interface UBinaryExpressionWithType : UExpression {
     override fun logString() = log("UBinaryExpressionWithType (${getExpressionType()?.name}, ${operationKind.name})", operand)
     override fun renderString() = "(${operand.renderString()}) ${operationKind.name} ${getExpressionType()?.name}"
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        operand.handleTraverse(callback)
+        if (visitor.visitBinaryExpressionWithType(this)) return
-        type.handleTraverse(callback)
+        operand.accept(visitor)
+        type.accept(visitor)
     }
 }

plugins/uast-common/src/org/jetbrains/uast/expressions/UBlockExpression.kt

@@ -15,10 +15,16 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UBlockExpression : UExpression {
     val expressions: List<UExpression>
 
-    override fun traverse(callback: UastCallback) = expressions.handleTraverseList(callback)
+    override fun accept(visitor: UastVisitor) {
+        if (visitor.visitBlockExpression(this)) return
+        expressions.acceptList(visitor)
+    }
+
     override fun logString() = log("UBlockExpression", expressions)
 
     override fun renderString() = buildString {

plugins/uast-common/src/org/jetbrains/uast/expressions/UCallExpression.kt

@@ -15,6 +15,8 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UCallExpression : UExpression, UResolvable {
     val functionReference: USimpleReferenceExpression?
 
@@ -36,12 +38,13 @@ interface UCallExpression : UExpression, UResolvable {
     override fun resolve(context: UastContext): UFunction?
     override fun resolveOrEmpty(context: UastContext): UFunction = resolve(context) ?: UFunctionNotResolved
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        functionReference?.handleTraverse(callback)
+        if (visitor.visitCallExpression(this)) return
-        classReference?.handleTraverse(callback)
+        functionReference?.accept(visitor)
-        functionNameElement?.handleTraverse(callback)
+        classReference?.accept(visitor)
-        valueArguments.handleTraverseList(callback)
+        functionNameElement?.accept(visitor)
-        typeArguments.handleTraverseList(callback)
+        valueArguments.acceptList(visitor)
+        typeArguments.acceptList(visitor)
     }
 
     override fun logString() = log("UFunctionCallExpression ($kind, argCount = $valueArgumentCount)", functionReference, valueArguments)

plugins/uast-common/src/org/jetbrains/uast/expressions/UCallableReferenceExpression.kt

@@ -15,11 +15,14 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UCallableReferenceExpression : UExpression {
     val qualifierType: UType
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        qualifierType.handleTraverse(callback)
+        if (visitor.visitCallableReferenceExpression(this)) return
+        qualifierType.accept(visitor)
     }
 
     override fun logString() = "UCallableReferenceExpression"

plugins/uast-common/src/org/jetbrains/uast/expressions/UClassLiteralExpression.kt

@@ -15,9 +15,15 @@
  */
 package org.jetbrains.uast
 
-interface UClassLiteralExpression : UExpression, LeafUElement {
+import org.jetbrains.uast.visitor.UastVisitor
+
+interface UClassLiteralExpression : UExpression {
     override fun logString() = "UClassLiteralExpression"
     override fun renderString() = getExpressionType()?.name ?: "::class"
 
     val type: UType
+
+    override fun accept(visitor: UastVisitor) {
+        visitor.visitClassLiteralExpression(this)
+    }
 }

plugins/uast-common/src/org/jetbrains/uast/expressions/UDeclarationsExpression.kt

@@ -15,14 +15,18 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UDeclarationsExpression : UExpression {
     val declarations: List<UElement>
 
     val variables: List<UVariable>
         get() = declarations.filterIsInstance<UVariable>()
 
-    override fun evaluate() = null
+    override fun accept(visitor: UastVisitor) {
-    override fun traverse(callback: UastCallback) = declarations.handleTraverseList(callback)
+        if (visitor.visitDeclarationsExpression(this)) return
+        declarations.acceptList(visitor)
+    }
 
     override fun renderString() = declarations.joinToString("\n") { it.renderString() }
     override fun logString() = log("UDeclarationsExpression", declarations)
@@ -31,6 +35,4 @@ interface UDeclarationsExpression : UExpression {
 class SimpleUDeclarationsExpression(
         override val parent: UElement,
         override val declarations: List<UElement>
-) : UDeclarationsExpression {
+) : UDeclarationsExpression
-    override fun evaluate() = null
-}

plugins/uast-common/src/org/jetbrains/uast/expressions/UExpressionList.kt

@@ -15,11 +15,15 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UExpressionList : UExpression {
     val expressions: List<UExpression>
 
-    override fun evaluate(): Any? = null
+    override fun accept(visitor: UastVisitor) {
-    override fun traverse(callback: UastCallback) = expressions.forEach { it.handleTraverse(callback) }
+        if (visitor.visitExpressionList(this)) return
+        expressions.acceptList(visitor)
+    }
 
     override fun logString() = log("UExpressionList", expressions)
     override fun renderString() = log("", expressions)

plugins/uast-common/src/org/jetbrains/uast/expressions/ULabeledExpression.kt

@@ -15,12 +15,15 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface ULabeledExpression : UExpression {
     val label: String
     val expression: UExpression
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        expression.handleTraverse(callback)
+        if (visitor.visitLabeledExpression(this)) return
+        expression.accept(visitor)
     }
 
     override fun evaluate() = expression.evaluate()

plugins/uast-common/src/org/jetbrains/uast/expressions/ULambdaExpression.kt

@@ -15,13 +15,16 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface ULambdaExpression : UExpression {
     val valueParameters: List<UVariable>
     val body: UExpression
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        valueParameters.handleTraverseList(callback)
+        if (visitor.visitLambdaExpression(this)) return
-        body.handleTraverse(callback)
+        valueParameters.acceptList(visitor)
+        body.accept(visitor)
     }
 
     override fun logString() = log("ULambdaExpression", valueParameters, body)

plugins/uast-common/src/org/jetbrains/uast/expressions/ULiteralExpression.kt

@@ -15,6 +15,8 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface ULiteralExpression : UExpression {
     val value: Any?
 
@@ -28,7 +30,10 @@ interface ULiteralExpression : UExpression {
 
     fun asString() = value?.toString() ?: ""
 
-    override fun traverse(callback: UastCallback) {}
+    override fun accept(visitor: UastVisitor) {
+        visitor.visitLiteralExpression(this)
+    }
+
     override fun logString() = "ULiteralExpression (${asString()})"
     override fun renderString() = asString()
 }

plugins/uast-common/src/org/jetbrains/uast/expressions/UNamedExpression.kt

@@ -15,14 +15,17 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 class UNamedExpression(
         override val name: String,
         override val parent: UElement
 ): UExpression, UNamed {
     lateinit var expression: UExpression
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        expression.handleTraverse(callback)
+        if (visitor.visitElement(this)) return
+        expression.accept(visitor)
     }
 
     override fun logString() = log("UNamedExpression ($name)", expression)

plugins/uast-common/src/org/jetbrains/uast/expressions/UObjectLiteralExpression.kt

@@ -15,11 +15,14 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UObjectLiteralExpression : UExpression {
     val declaration: UClass
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        declaration.handleTraverse(callback)
+        if (visitor.visitObjectLiteralExpression(this)) return
+        declaration.accept(visitor)
     }
 
     override fun logString() = log("UObjectLiteralExpression", declaration)

plugins/uast-common/src/org/jetbrains/uast/expressions/UParenthesizedExpression.kt

@@ -15,11 +15,14 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UParenthesizedExpression : UExpression {
     val expression: UExpression
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        expression.handleTraverse(callback)
+        if (visitor.visitParenthesizedExpression(this)) return
+        expression.accept(visitor)
     }
 
     override fun evaluate() = expression.evaluate()

plugins/uast-common/src/org/jetbrains/uast/expressions/UQualifiedExpression.kt

@@ -15,6 +15,8 @@
  */
 package org.jetbrains.uast
 
+import org.jetbrains.uast.visitor.UastVisitor
+
 interface UQualifiedExpression : UExpression, UResolvable {
     val receiver: UExpression
     val selector: UExpression
@@ -24,9 +26,10 @@ interface UQualifiedExpression : UExpression, UResolvable {
 
     override fun renderString() = receiver.renderString() + accessType.name + selector.renderString()
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        receiver.handleTraverse(callback)
+        if (visitor.visitQualifiedExpression(this)) return
-        selector.handleTraverse(callback)
+        receiver.accept(visitor)
+        selector.accept(visitor)
     }
 
     override fun logString() = log("UQualifiedExpression", receiver, selector)

plugins/uast-common/src/org/jetbrains/uast/expressions/USimpleReferenceExpression.kt

@@ -15,9 +15,15 @@
  */
 package org.jetbrains.uast
 
-interface USimpleReferenceExpression : UExpression, UResolvable, LeafUElement {
+import org.jetbrains.uast.visitor.UastVisitor
+
+interface USimpleReferenceExpression : UExpression, UResolvable {
     val identifier: String
 
+    override fun accept(visitor: UastVisitor) {
+        visitor.visitSimpleReferenceExpression(this)
+    }
+
     override fun logString() = "USimpleReferenceExpression ($identifier)"
     override fun renderString() = identifier
 }

plugins/uast-common/src/org/jetbrains/uast/expressions/USuperExpression.kt

@@ -15,7 +15,13 @@
  */
 package org.jetbrains.uast
 
-interface USuperExpression : UExpression, LeafUElement {
+import org.jetbrains.uast.visitor.UastVisitor
+
+interface USuperExpression : UExpression {
     override fun logString() = "USuperExpression"
     override fun renderString() = "super"
+
+    override fun accept(visitor: UastVisitor) {
+        visitor.visitSuperExpression(this)
+    }
 }

plugins/uast-common/src/org/jetbrains/uast/expressions/UThisExpression.kt

@@ -15,7 +15,13 @@
  */
 package org.jetbrains.uast
 
-interface UThisExpression : UExpression, LeafUElement {
+import org.jetbrains.uast.visitor.UastVisitor
+
+interface UThisExpression : UExpression {
     override fun logString() = "UThisExpression"
     override fun renderString() = "this"
+
+    override fun accept(visitor: UastVisitor) {
+        visitor.visitThisExpression(this)
+    }
 }

plugins/uast-common/src/org/jetbrains/uast/expressions/UUnaryExpression.kt

@@ -16,24 +16,38 @@
 package org.jetbrains.uast
 
 import org.jetbrains.uast.kinds.UastOperator
+import org.jetbrains.uast.visitor.UastVisitor
 
 interface UUnaryExpression : UExpression {
     val operand: UExpression
     val operator: UastOperator
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        operand.handleTraverse(callback)
+        if (visitor.visitUnaryExpression(this)) return
+        operand.accept(visitor)
     }
 }
 
 interface UPrefixExpression : UUnaryExpression {
     override val operator: UastPrefixOperator
+
+    override fun accept(visitor: UastVisitor) {
+        if (visitor.visitPrefixExpression(this)) return
+        operand.accept(visitor)
+    }
+
     override fun logString() = log("UPrefixExpression (${operator.text})", operand)
     override fun renderString() = operator.text + operand.renderString()
 }
 
 interface UPostfixExpression : UUnaryExpression {
     override val operator: UastPostfixOperator
+
+    override fun accept(visitor: UastVisitor) {
+        if (visitor.visitPostfixExpression(this)) return
+        operand.accept(visitor)
+    }
+
     override fun logString() = log("UPostfixExpression (${operator.text})", operand)
     override fun renderString() = operand.renderString() + operator.text
 }

plugins/uast-java/src/org/jetbrains/uast/java/expressions/JavaDumbUElement.kt

@@ -16,14 +16,13 @@
 package org.jetbrains.uast.java
 
 import com.intellij.psi.PsiElement
-import org.jetbrains.uast.LeafUElement
 import org.jetbrains.uast.UElement
 import org.jetbrains.uast.psi.PsiElementBacked
 
 class JavaDumbUElement(
         override val psi: PsiElement?,
         override val parent: UElement
-) : JavaAbstractUElement(), UElement, PsiElementBacked, LeafUElement {
+) : JavaAbstractUElement(), UElement, PsiElementBacked {
     override fun logString() = "JavaDumbUElement"
     override fun renderString() = "<stub@$psi>"
 }

plugins/uast-java/src/org/jetbrains/uast/java/expressions/JavaUInstanceCheckExpression.kt

@@ -30,6 +30,4 @@ class JavaUInstanceCheckExpression(
 
     override val operationKind: UastBinaryExpressionWithTypeKind.InstanceCheck
         get() = UastBinaryExpressionWithTypeKind.INSTANCE_CHECK
-
-    override fun evaluate() = null
 }

plugins/uast-java/src/org/jetbrains/uast/java/expressions/JavaULiteralExpression.kt

@@ -24,7 +24,7 @@ class JavaULiteralExpression(
         override val psi: PsiLiteralExpression,
         override val parent: UElement
 ) : JavaAbstractUElement(), ULiteralExpression, PsiElementBacked, JavaUElementWithType {
-    override val asString by lz { psi.text }
+    override fun asString() = psi.text
     override fun evaluate() = psi.value
     override val value by lz { evaluate() }
 

plugins/uast-java/src/org/jetbrains/uast/java/expressions/UnknownJavaExpression.kt

@@ -22,6 +22,6 @@ import org.jetbrains.uast.psi.PsiElementBacked
 class UnknownJavaExpression(
         override val psi: PsiElement,
         override val parent: UElement
-) : UExpression, PsiElementBacked, LeafUElement {
+) : UExpression, PsiElementBacked {
     override fun logString() = "[!] UnknownJavaExpression ($psi)"
 }

plugins/uast-java/src/org/jetbrains/uast/java/expressions/javaUCallExpressions.kt

@@ -144,7 +144,6 @@ class JavaArrayInitializerUCallExpression(
         get() = JavaUastCallKinds.ARRAY_INITIALIZER
 
     override fun resolve(context: UastContext) = null
-    override fun evaluate() = null
 }
 
 class JavaAnnotationArrayInitializerUCallExpression(

plugins/uast-kotlin/src/org/jetbrains/kotlin/uast/declarations/KotlinDumbUElement.kt

@@ -17,14 +17,13 @@
 package org.jetbrains.kotlin.uast
 
 import com.intellij.psi.PsiElement
-import org.jetbrains.uast.LeafUElement
 import org.jetbrains.uast.UElement
 import org.jetbrains.uast.psi.PsiElementBacked
 
 class KotlinDumbUElement(
         override val psi: PsiElement?,
         override val parent: UElement
-) : KotlinAbstractUElement(), UElement, PsiElementBacked, LeafUElement {
+) : KotlinAbstractUElement(), UElement, PsiElementBacked {
     override fun logString() = "KotlinDumbUElement"
     override fun renderString() = "<stub@$psi>"
 }

plugins/uast-kotlin/src/org/jetbrains/kotlin/uast/declarations/KotlinUAnnotationList.kt

@@ -19,9 +19,9 @@ package org.jetbrains.kotlin.uast
 import org.jetbrains.kotlin.psi.KtAnnotation
 import org.jetbrains.uast.UAnnotation
 import org.jetbrains.uast.UElement
-import org.jetbrains.uast.UastCallback
+import org.jetbrains.uast.acceptList
-import org.jetbrains.uast.handleTraverseList
 import org.jetbrains.uast.psi.PsiElementBacked
+import org.jetbrains.uast.visitor.UastVisitor
 
 class KotlinUAnnotationList(
         override val psi: KtAnnotation,
@@ -31,7 +31,9 @@ class KotlinUAnnotationList(
 
     override fun logString() = "KotlinUAnnotationList"
     override fun renderString() = annotations.joinToString(" ") { it.renderString() }
-    override fun traverse(callback: UastCallback) {
+
-        annotations.handleTraverseList(callback)
+    override fun accept(visitor: UastVisitor) {
+        if (visitor.visitElement(this)) return
+        annotations.acceptList(visitor)
     }
 }

plugins/uast-kotlin/src/org/jetbrains/kotlin/uast/expressions/KotlinUDeclarationsExpression.kt

@@ -20,6 +20,4 @@ import org.jetbrains.kotlin.uast.KotlinAbstractUElement
 class KotlinUDeclarationsExpression(override val parent: UElement) : KotlinAbstractUElement(), UDeclarationsExpression {
     override lateinit var declarations: List<UElement>
         internal set
-
-    override fun evaluate() = null
 }

plugins/uast-kotlin/src/org/jetbrains/kotlin/uast/expressions/KotlinULiteralExpression.kt

@@ -40,8 +40,7 @@ class KotlinStringULiteralExpression(
         override val parent: UElement
 ) : KotlinAbstractUElement(), ULiteralExpression, PsiElementBacked, KotlinUElementWithType, KotlinEvaluatableUElement {
     override val isNull = false
-    override val asString: String
+    override fun asString(): String = '"' + psi.text + '"'
-        get() = '"' + psi.text + '"'
 
     override val value: String
         get() = psi.text

plugins/uast-kotlin/src/org/jetbrains/kotlin/uast/expressions/KotlinUSwitchExpression.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.uast.kinds.KotlinSpecialExpressionKinds
 import org.jetbrains.uast.*
 import org.jetbrains.uast.psi.PsiElementBacked
+import org.jetbrains.uast.visitor.UastVisitor
 
 class KotlinUSwitchExpression(
         override val psi: KtWhenExpression,
@@ -101,7 +102,8 @@ class KotlinUSwitchEntry(
 
     override fun logString() = log("KotlinUSwitchEntry", expression)
 
-    override fun traverse(callback: UastCallback) {
+    override fun accept(visitor: UastVisitor) {
-        expression.handleTraverse(callback)
+        if (visitor.visitElement(this)) return
+        expression.accept(visitor)
     }
 }

plugins/uast-kotlin/src/org/jetbrains/kotlin/uast/expressions/UnknownKotlinExpression.kt

@@ -19,13 +19,11 @@ package org.jetbrains.kotlin.uast
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.uast.UElement
 import org.jetbrains.uast.UExpression
-import org.jetbrains.uast.UastCallback
 import org.jetbrains.uast.psi.PsiElementBacked
 
 class UnknownKotlinExpression(
         override val psi: KtExpression,
         override val parent: UElement
 ) : KotlinAbstractUElement(), UExpression, PsiElementBacked {
-    override fun traverse(callback: UastCallback) {}
     override fun logString() = "[!] UnknownKotlinExpression ($psi)"
 }

plugins/uast-kotlin/test/org.jetbrains.kotlin.uast/AbstractKotlinUastStructureTest.kt

@@ -19,6 +19,8 @@ package org.jetbrains.kotlin.uast
 import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
 import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor
 import org.jetbrains.kotlin.test.KotlinTestUtils
+import org.jetbrains.uast.UElement
+import org.jetbrains.uast.visitor.UastVisitor
 import java.io.File
 
 abstract class AbstractKotlinUastStructureTest : KotlinLightCodeInsightFixtureTestCase() {
@@ -28,6 +30,7 @@ abstract class AbstractKotlinUastStructureTest : KotlinLightCodeInsightFixtureTe
 
         val logFile = File(File(testDataPath, "log"), "$testName.txt")
         val renderFile = File(File(testDataPath, "render"), "$testName.txt")
+        val treeFile = File(File(testDataPath, "tree"), "$testName.txt")
 
         val psiFile = myFixture.file
         val uElement = KotlinUastLanguagePlugin.converter.convertWithParent(psiFile) ?: error("UFile was not created")
@@ -42,6 +45,24 @@ abstract class AbstractKotlinUastStructureTest : KotlinLightCodeInsightFixtureTe
             throw e
         }
         KotlinTestUtils.assertEqualsToFile(renderFile, renderActual)
+        KotlinTestUtils.assertEqualsToFile(treeFile, genTree(uElement))
+    }
+
+    private fun genTree(node: UElement): String {
+        val builder = StringBuilder()
+        val visitor = object : UastVisitor() {
+            private tailrec fun getParentCount(node: UElement): Int {
+                val parent = node.parent ?: return 0
+                return getParentCount(parent)
+            }
+
+            override fun visitElement(node: UElement): Boolean {
+                builder.appendln("    ".repeat(getParentCount(node)) + node.javaClass.name)
+                return super.visitElement(node)
+            }
+        }
+        visitor.visitElement(node)
+        return builder.toString()
     }
 
     override fun getTestDataPath() = "plugins/uast-kotlin/testData"