Fix AIOOBE when completion started on empty file
[idea/community.git] / plugins / stats-collector / src / com / intellij / stats / ngram / NGram.kt
1 package com.intellij.stats.ngram
2
3 import com.intellij.psi.PsiElement
4 import com.intellij.psi.PsiFile
5 import com.intellij.util.io.DataInputOutputUtil
6 import com.intellij.util.io.KeyDescriptor
7 import java.io.DataInput
8 import java.io.DataOutput
9
10 data class NGram(val elements: List<String>) {
11
12     fun append(token: String): NGram {
13         return NGram(listOf(*elements.toTypedArray(), token))
14     }
15
16     fun dropHead(): NGram {
17         return NGram(elements.subList(1, elements.size))
18     }
19
20     companion object {
21         val INVALID = NGram(emptyList())
22
23         val N = 5
24
25         fun processFile(psiFile: PsiFile, content: CharSequence): HashMap<NGram, Int> {
26             val result = HashMap<NGram, Int>()
27             val elements = TreeTraversal.getElements(psiFile)
28             if (elements.size > N) {
29                 for (i in N until elements.size) {
30                     if (NGramElementProvider.shouldIndex(elements[i], content)) {
31                         val nGramElements = ArrayList<String>()
32                         for (j in 1..N) {
33                             nGramElements.add(NGramElementProvider.getElementRepresentation(elements[i - j]))
34                         }
35                         nGramElements.add(elements[i].text)
36                         for (j in 0..N) {
37                             val ngram = NGram(nGramElements.subList(j, nGramElements.size))
38                             val oldValue = result.putIfAbsent(ngram, 1)
39                             oldValue?.let {
40                                 result.put(ngram, oldValue + 1)
41                             }
42                         }
43                     }
44                 }
45             }
46             return result
47         }
48
49         fun getNGramForElement(element: PsiElement) : NGram {
50             val elements = TreeTraversal.getElements(element.containingFile ?: return NGram.INVALID)
51             val index = elements.indexOf(element.parent.node)
52             if (index == -1 || index < NGram.N) {
53                 return NGram.INVALID
54             }
55             val nGramElements = ArrayList<String>()
56             for (i in 1..NGram.N) {
57                 nGramElements.add(elements[index - i].elementType.toString())
58             }
59             return NGram(nGramElements)
60         }
61     }
62 }
63
64 fun DataOutput.writeINT(x : Int) = DataInputOutputUtil.writeINT(this, x)
65 fun DataInput.readINT() : Int = DataInputOutputUtil.readINT(this)
66
67 object NGramKeyDescriptor: KeyDescriptor<NGram> {
68     override fun save(out: DataOutput, nGram: NGram?) {
69         val instance = NGramEnumeratingService.getInstance()
70         out.writeINT(nGram!!.elements.size)
71         nGram.elements.forEach { out.writeINT(instance.enumerateString(it)) }
72     }
73
74     override fun read(`in`: DataInput): NGram {
75         val instance = NGramEnumeratingService.getInstance()
76         val size = `in`.readINT()
77         return NGram((1..size).map { instance.valueOf(`in`.readINT()) })
78     }
79
80     override fun isEqual(p0: NGram?, p1: NGram?): Boolean {
81         return p0 == p1
82     }
83
84     override fun getHashCode(p0: NGram?): Int {
85         return p0!!.hashCode()
86     }
87 }