diff options
Diffstat (limited to 'lua/telescope/utils.lua')
| -rw-r--r-- | lua/telescope/utils.lua | 122 |
1 files changed, 122 insertions, 0 deletions
diff --git a/lua/telescope/utils.lua b/lua/telescope/utils.lua index 196453a..e758ed0 100644 --- a/lua/telescope/utils.lua +++ b/lua/telescope/utils.lua @@ -11,4 +11,126 @@ utils.reversed_ipairs = function(t) return reversedipairsiter, t, #t + 1 end +utils.default_table_mt = { + __index = function(t, k) + local obj = {} + rawset(t, k, obj) + return obj + end +} + +local NGram = {} +NGram.__index = NGram + +function NGram:new(opts) + -- TODO: Add padding + opts = opts or {} + return setmetatable({ + N = opts.N or 2, + split = opts.split or "/", + _depth = 5, + _grams = setmetatable({}, utils.default_table_mt) + }, self) +end + +local min = math.min + +function NGram:_split(word) + local word_len = #word + + local result = {} + for i = 1, word_len - 1 do + -- for j = i + (self.N - 1), min(i + self._depth - 1, word_len) do + -- table.insert(result, string.sub(word, i, j)) + -- end + table.insert(result, string.sub(word, i, i + self.N - 1)) + end + + return result +end + +-- local function pairsByKeys (t, f) +-- local a = {} +-- for n in pairs(t) do table.insert(a, n) end +-- table.sort(a, f) +-- local i = 0 -- iterator variable +-- local iter = function () -- iterator function +-- i = i + 1 +-- if a[i] == nil then return nil +-- else return a[i], t[a[i]] +-- end +-- end +-- return iter +-- end + +function NGram:add(word) + local split_word = self:_split(word) + + for _, k in ipairs(split_word) do + local counts = self._grams[k] + if counts[word] == nil then + counts[word] = 0 + end + + counts[word] = counts[word] + 1 + end +end + +function NGram:_items_sharing_ngrams(query) + local split_query = self:_split(query) + + -- Matched string to number of N-grams shared with the query string. + local shared = {} + + local remaining = {} + + for _, ngram in ipairs(split_query) do + remaining = {} + for match, count in pairs(self._grams[ngram] or {}) do + remaining[match] = remaining[match] or count + + if remaining[match] > 0 then + remaining[match] = remaining[match] - 1 + shared[match] = (shared[match] or 0) + 1 + end + end + end + + return shared +end + +function NGram:search(query, show_values) + local sharing_ngrams = self:_items_sharing_ngrams(query) + + local results = {} + for name, count in pairs(sharing_ngrams) do + local allgrams = #query + #name - (2 * self.N) - count + 2 + table.insert(results, {name, count / allgrams}) + end + + table.sort(results, function(left, right) + return left[2] > right[2] + end) + + if not show_values then + for k, v in ipairs(results) do + results[k] = v[1] + end + end + + return results +end + +function NGram:find(query) + return self:search(query)[1] +end + +function NGram:score(query) + return (self:search(query, true)[1] or {})[2] or 0 +end + +utils.new_ngram = function() + return NGram:new() +end + return utils |
