package com.intellij.tokenindex;

import com.intellij.lang.ASTNode;
import com.intellij.lang.javascript.JSTokenTypes;
import com.intellij.lang.javascript.psi.*;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiWhiteSpace;
import com.intellij.psi.impl.source.tree.LeafElement;
import org.jetbrains.annotations.NotNull;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

/**
 * @author Eugene.Kudelevsky
 */
public class JavaScriptTokenizer implements Tokenizer {
  private static final byte EXPRESSION_TYPE = 0;
  private static final byte IDENTIFIER_TYPE = 1;
  private static final byte PARAM_LIST_TYPE = 2;

  @NotNull
  public List<Token> tokenize(Collection<? extends PsiElement> roots) {
    MyVisitor visitor = new MyVisitor();
    for (PsiElement root : roots) {
      root.accept(visitor);
    }
    return visitor.getTokens();
  }

  private static class MyVisitor extends JSRecursiveElementVisitor {
    private final List<Token> myTokens = new ArrayList<Token>();

    @Override
    public void visitElement(PsiElement element) {
      if (element instanceof JSExpression) {
        myTokens.add(new AnonymToken(EXPRESSION_TYPE, element.getTextOffset()));
      }
      else {
        super.visitElement(element);
        if (element instanceof PsiWhiteSpace || !(element instanceof LeafElement)) {
          return;
        }
        visitLeafElement(element);
      }
    }

    @Override
    public void visitJSParameterList(JSParameterList node) {
      myTokens.add(new AnonymToken(PARAM_LIST_TYPE, node.getTextOffset()));
    }

    @Override
    public void visitJSBlock(JSBlockStatement node) {
      PsiElement parent = node.getParent();
      if (parent instanceof JSIfStatement || parent instanceof JSLoopStatement) {
        for (JSStatement statement : node.getStatements()) {
          statement.accept(this);
        }
      }
      else {
        super.visitJSBlock(node);
      }
    }

    private void visitLeafElement(PsiElement element) {
      if (element.getTextLength() == 0) {
        return;
      }
      ASTNode node = element.getNode();
      if (node != null && node.getElementType() == JSTokenTypes.IDENTIFIER) {
        myTokens.add(new AnonymToken(IDENTIFIER_TYPE, element.getTextOffset()));
      }
      else {
        myTokens.add(new TextToken(element.getText().hashCode(), element.getTextOffset()));
      }
    }

    public List<Token> getTokens() {
      return myTokens;
    }
  }
}

@@ -279,24 +279,78 @@
 eafElement;%0A
+import com.intellij.psi.impl.source.tree.TreeElement;%0A
 import org.j
@@ -1066,16 +1066,53 @@
 oken%3E();
+%0A    private int myParentOffset = -1;
 %0A%0A    @O
@@ -1181,33 +1181,340 @@
    i
-f (element instanceof JSE
+nt temp = myParentOffset;%0A      myParentOffset = getTextOffset(element);%0A      super.visitElement(element);%0A      myParentOffset = temp;%0A      if (element instanceof LeafElement && !(element instanceof PsiWhiteSpace)) %7B%0A        visitLeafElement(element);%0A      %7D%0A    %7D%0A%0A    @Override%0A    public void visitJSExpression(JSExpression e
 xpre
@@ -1518,26 +1518,24 @@
 pression) %7B%0A
-  
       myToke
@@ -1570,32 +1570,24 @@
 SSION_TYPE, 
-element.
 getTextOffse
@@ -1580,32 +1580,42 @@
 , getTextOffset(
+expression
 )));%0A      %7D%0A   
@@ -1611,119 +1611,244 @@
     
-  
 %7D
+%0A
 %0A    
-  else %7B%0A        super.visitElement(element);%0A        if (element instanceof PsiWhiteSpace %7C%7C !(element
+private int getTextOffset(@NotNull PsiElement element) %7B%0A      // performance hint: TreeElement#getTextOffset() works recursively upwards%0A      if (myParentOffset %3E= 0) %7B%0A        ASTNode node = element.getNode();%0A        if (node
  ins
@@ -1859,20 +1859,20 @@
 eof 
-Leaf
+Tree
 Element)
 ) %7B%0A
@@ -1867,17 +1867,16 @@
 Element)
-)
  %7B%0A     
@@ -1890,62 +1890,127 @@
 turn
-;%0A        %7D%0A        visitLeafElement(element);%0A      %7D
+ myParentOffset + ((TreeElement)node).getStartOffsetInParent();%0A        %7D%0A      %7D%0A      return element.getTextOffset();
 %0A   
@@ -2140,21 +2140,16 @@
 T_TYPE, 
-node.
 getTextO
@@ -2146,32 +2146,36 @@
 , getTextOffset(
+node
 )));%0A    %7D%0A%0A    
@@ -2833,24 +2833,16 @@
 R_TYPE, 
-element.
 getTextO
@@ -2839,32 +2839,39 @@
 , getTextOffset(
+element
 )));%0A      %7D%0A   
@@ -2945,24 +2945,16 @@
 Code(), 
-element.
 getTextO
@@ -2959,16 +2959,23 @@
 tOffset(
+element
 )));%0A   

clean up, performance optimization