Work on ngram sorter

6 files changed, 365 insertions, 29 deletions

diff --git a/lua/telescope/finders.lua b/lua/telescope/finders.lua
index 518dfc1..e5a601e 100644
--- a/lua/telescope/finders.lua
+++ b/lua/telescope/finders.lua
@@ -1,4 +1,5 @@
 local a = vim.api
+local log = require('telescope.log')
 
 local finders = {}
 
@@ -29,7 +30,7 @@ function Finder:new(opts)
   --    ...
   return setmetatable({
     fn_command = opts.fn_command,
-    responsive = opts.responsive,
+    static = opts.static,
     state = {},
     job_id = -1,
   }, Finder)
@@ -44,11 +45,28 @@ end
 -- process_search
 -- do_your_job
 -- process_plz
-function Finder:_find(prompt, process_result)
+function Finder:_find(prompt, process_result, process_complete)
   if (self.state.job_id or 0) > 0 then
     vim.fn.jobstop(self.job_id)
   end
 
+  log.info("Finding...")
+  if self.static and self.done then
+    log.info("Using previous results")
+    for _, v in ipairs(self._cached_lines) do
+      process_result(v)
+    end
+
+    process_complete()
+    return
+  end
+
+  if self.static then
+    self._cached_lines = {}
+  end
+
+  self.done = false
+
   -- TODO: How to just literally pass a list...
   -- TODO: How to configure what should happen here
   -- TODO: How to run this over and over?
@@ -57,9 +75,21 @@ function Finder:_find(prompt, process_result)
 
     on_stdout = function(_, data, _)
       for _, line in ipairs(data) do
-        process_result(line)
+        if vim.trim(line) ~= "" then
+          process_result(line)
+
+          if self.static then
+            table.insert(self._cached_lines, line)
+          end
+        end
       end
-    end
+    end,
+
+    on_exit = function()
+      self.done = true
+
+      process_complete()
+    end,
   })
 end
 
diff --git a/lua/telescope/log.lua b/lua/telescope/log.lua
new file mode 100644
index 0000000..fb3d712
--- /dev/null
+++ b/lua/telescope/log.lua
@@ -0,0 +1,83 @@
+-- https://raw.githubusercontent.com/rxi/log.lua/master/log.lua
+-- log.lua
+--
+-- Copyright (c) 2016 rxi
+--
+-- This library is free software; you can redistribute it and/or modify it
+-- under the terms of the MIT license. See LICENSE for details.
+--
+
+local log = { _version = "0.1.0" }
+
+log.usecolor = true
+log.outfile = vim.fn.stdpath('data') .. '/telescope.log'
+log.console = false
+log.level = "trace"
+
+
+local modes = {
+  { name = "trace", color = "\27[34m", },
+  { name = "debug", color = "\27[36m", },
+  { name = "info",  color = "\27[32m", },
+  { name = "warn",  color = "\27[33m", },
+  { name = "error", color = "\27[31m", },
+  { name = "fatal", color = "\27[35m", },
+}
+
+
+local levels = {}
+for i, v in ipairs(modes) do
+  levels[v.name] = i
+end
+
+
+local round = function(x, increment)
+  increment = increment or 1
+  x = x / increment
+  return (x > 0 and math.floor(x + .5) or math.ceil(x - .5)) * increment
+end
+
+for i, x in ipairs(modes) do
+  local nameupper = x.name:upper()
+  log[x.name] = function(...)
+    -- Return early if we're below the log level
+    if i < levels[log.level] then
+      return
+    end
+
+    local passed = {...}
+    local fmt = table.remove(passed, 1)
+    local inspected = {}
+    for _, v in ipairs(passed) do
+      table.insert(inspected, vim.inspect(v))
+    end
+    local msg = string.format(fmt, unpack(inspected))
+    local info = debug.getinfo(2, "Sl")
+    local lineinfo = info.short_src .. ":" .. info.currentline
+
+    -- Output to console
+    if log.console then
+      print(string.format("%s[%-6s%s]%s %s: %s",
+                          log.usecolor and x.color or "",
+                          nameupper,
+                          os.date("%H:%M:%S"),
+                          log.usecolor and "\27[0m" or "",
+                          lineinfo,
+                          msg))
+    end
+
+    -- Output to log file
+    if log.outfile then
+      local fp = io.open(log.outfile, "a")
+      local str = string.format("[%-6s%s] %s: %s\n",
+                                nameupper, os.date(), lineinfo, msg)
+      fp:write(str)
+      fp:close()
+    end
+
+  end
+end
+
+log.info("Logger Succesfully Loaded")
+
+return log
diff --git a/lua/telescope/pickers.lua b/lua/telescope/pickers.lua
index 5b4485d..db5ec7d 100644
--- a/lua/telescope/pickers.lua
+++ b/lua/telescope/pickers.lua
@@ -1,6 +1,12 @@
 local a = vim.api
+local fun = require('fun')
 local popup = require('popup')
 
+
+local zip = fun.zip
+local tomap = fun.tomap
+
+local log = require('telescope.log')
 local mappings = require('telescope.mappings')
 local state = require('telescope.state')
 local utils = require('telescope.utils')
@@ -119,10 +125,10 @@ function Picker:find(opts)
   -- vim.fn.prompt_setprompt(prompt_bufnr, prompt_string)
 
   -- First thing we want to do is set all the lines to blank.
-  local max_results = popup_opts.results.height
+  self.max_results = popup_opts.results.height - 1
   local initial_lines = {}
-  for _ = 1, max_results do table.insert(initial_lines, "") end
-  vim.api.nvim_buf_set_lines(results_bufnr, 0, max_results, false, initial_lines)
+  for _ = 1, self.max_results do table.insert(initial_lines, "") end
+  vim.api.nvim_buf_set_lines(results_bufnr, 0, self.max_results, false, initial_lines)
 
   local on_lines = function(_, _, _, first_line, last_line)
     local prompt = vim.api.nvim_buf_get_lines(prompt_bufnr, first_line, last_line, false)[1]
@@ -145,11 +151,22 @@ function Picker:find(opts)
     -- TODO: We need to handle huge lists in a good way, cause currently we'll just put too much stuff in the buffer
     -- TODO: Stop having things crash if we have an error.
 
-    local replace_line = function(row, line)
+    local replace_line = function(score, row, line)
+      log.trace("Replacing @ %s w/ text '%s' (%s)", row, line, score)
       vim.api.nvim_buf_set_lines(results_bufnr, row, row + 1, false, {line})
     end
 
-    finder(prompt, function(line)
+    local insert_line = function(score, row, line)
+      log.trace("Inserting @ %s w/ text '%s' (%s)", row, line, score)
+      vim.api.nvim_buf_set_lines(results_bufnr, row, row, false, {line})
+    end
+
+
+    local process_result = function(line)
+      if vim.trim(line) == "" then
+        return
+      end
+
       if sorter then
         local sort_score = sorter:score(prompt, line)
         if sort_score == -1 then
@@ -161,7 +178,7 @@ function Picker:find(opts)
         for row, row_score in utils.reversed_ipairs(self.line_scores) do
           if row_score > sort_score then
             -- Insert line at row
-            replace_line(max_results - row, line)
+            insert_line(sort_score, self.max_results - row, line)
 
             -- Insert current score in the table
             table.insert(self.line_scores, row + 1, sort_score)
@@ -171,7 +188,7 @@ function Picker:find(opts)
           end
 
           -- Don't keep inserting stuff
-          if row > max_results then
+          if row > self.max_results then
             return
           end
         end
@@ -179,19 +196,35 @@ function Picker:find(opts)
         -- Worst score so far, so add to end
 
         -- example: 5 max results, 8
-        local worst_line = max_results - #self.line_scores
-        replace_line(worst_line, line)
+        local worst_line = self.max_results - #self.line_scores
+        replace_line(sort_score, worst_line, line)
         table.insert(self.line_scores, sort_score)
       else
         -- Just always append to the end of the buffer if this is all you got.
         vim.api.nvim_buf_set_lines(results_bufnr, -1, -1, false, {line})
       end
+    end
+
+    local process_complete = function()
+      local worst_line = self.max_results - #self.line_scores
+      local empty_lines = {}
+      for _ = 1, worst_line do table.insert(empty_lines, "") end
+      vim.api.nvim_buf_set_lines(results_bufnr, 0, worst_line, false, empty_lines)
+
+      log.info("Worst Line after process_complete: %s", worst_line)
+      log.trace("%s", tomap(zip(
+        a.nvim_buf_get_lines(results_bufnr, worst_line, self.max_results, false),
+        self.line_scores
+      )))
+    end
+
+    pcall(function()
+      return finder(prompt, process_result, process_complete)
     end)
-    -- local results = finder:get_results(results_win, results_bufnr, line)
   end
 
   -- Call this once to pre-populate if it makes sense
-  vim.schedule_wrap(on_lines(nil, nil, nil, 0, 1))
+  -- vim.schedule_wrap(on_lines(nil, nil, nil, 0, 1))
 
   -- Register attach
   vim.api.nvim_buf_attach(prompt_bufnr, true, {
@@ -233,16 +266,12 @@ function Picker:find(opts)
     finder = finder,
   })
 
-  -- print(vim.inspect(state.get_status(prompt_bufnr)))
   mappings.set_keymap(prompt_bufnr, results_bufnr)
 
   vim.cmd [[startinsert]]
 end
 
 function Picker:close_windows(status)
-  -- vim.fn['popup#close_win'](state.prompt_win)
-  -- vim.fn['popup#close_win'](state.results_win)
-  -- vim.fn['popup#close_win'](state.preview_win)
   local prompt_win = status.prompt_win
   local results_win = status.results_win
   local preview_win = status.preview_win
@@ -252,12 +281,13 @@ function Picker:close_windows(status)
   local preview_border_win = status.preview_border_win
 
   local function del_win(name, win_id, force)
-    -- local file = io.open("/home/tj/test.txt", "a")
-    -- file:write(string.format("Closing.... %s %s\n", name, win_id))
-    local ok = pcall(vim.api.nvim_win_close, win_id, force)
-    -- file:write(string.format("OK: %s\n", ok))
-    -- file:write("...Done\n\n")
-    -- file:close()
+    if not vim.api.nvim_win_is_valid(win_id) then
+      return
+    end
+
+    if not pcall(vim.api.nvim_win_close, win_id, force) then
+      log.trace("Unable to close window: %s/%s", name, win_id)
+    end
   end
 
   del_win("prompt_win", prompt_win, true)
@@ -283,7 +313,7 @@ end
 local ns_telescope_selection = a.nvim_create_namespace('telescope_selection')
 
 function Picker:get_selection()
-  return self.selection or #(self.line_scores or {})
+  return self.selection or self.max_results
 end
 
 function Picker:move_selection(change)
@@ -291,6 +321,12 @@ function Picker:move_selection(change)
 end
 
 function Picker:set_selection(row)
+  if row > self.max_results then
+    row = self.max_results
+  elseif row < 1 then
+    row = 1
+  end
+
   local status = state.get_status(self.prompt_bufnr)
 
   a.nvim_buf_clear_namespace(status.results_bufnr, ns_telescope_selection, 0, -1)
diff --git a/lua/telescope/previewers.lua b/lua/telescope/previewers.lua
index 5b686b7..31ea364 100644
--- a/lua/telescope/previewers.lua
+++ b/lua/telescope/previewers.lua
@@ -1,3 +1,5 @@
+local log = require('telescope.log')
+
 local previewers = {}
 
 local Previewer = {}
@@ -21,12 +23,14 @@ end
 
 previewers.vim_buffer = previewers.new {
   preview_fn = function(preview_win, preview_bufnr, results_bufnr, row)
-    assert(preview_bufnr)
-
     local line = vim.api.nvim_buf_get_lines(results_bufnr, row, row + 1, false)[1]
+    if line == nil then
+      return
+    end
     local file_name = vim.split(line, ":")[1]
 
-    -- print(file_name)
+    log.info("Previewing File: %s", file_name)
+
     -- vim.fn.termopen(
     --   string.format("bat --color=always --style=grid %s"),
     -- vim.fn.fnamemodify(file_name, ":p")
@@ -45,6 +49,39 @@ previewers.vim_buffer = previewers.new {
 }
 
 
+previewers.vim_buffer_or_bat = previewers.new {
+  preview_fn = function(preview_win, preview_bufnr, results_bufnr, row)
+    local line = vim.api.nvim_buf_get_lines(results_bufnr, row, row + 1, false)[1]
+    if line == nil then
+      return
+    end
+    local file_name = vim.split(line, ":")[1]
+
+    log.info("Previewing File: %s", file_name)
+
+    -- vim.fn.termopen(
+    --   string.format("bat --color=always --style=grid %s"),
+    -- vim.fn.fnamemodify(file_name, ":p")
+    local bufnr = vim.fn.bufadd(file_name)
+
+    if vim.api.nvim_buf_is_loaded(bufnr) then
+      vim.fn.bufload(bufnr)
+
+      -- TODO: We should probably call something like this because we're not always getting highlight and all that stuff.
+      -- api.nvim_command('doautocmd filetypedetect BufRead ' .. vim.fn.fnameescape(filename))
+      vim.api.nvim_win_set_buf(preview_win, bufnr)
+      vim.api.nvim_win_set_option(preview_win, 'wrap', false)
+      vim.api.nvim_win_set_option(preview_win, 'winhl', 'Normal:Normal')
+      -- vim.api.nvim_win_set_option(preview_win, 'winblend', 20)
+      vim.api.nvim_win_set_option(preview_win, 'signcolumn', 'no')
+      vim.api.nvim_win_set_option(preview_win, 'foldlevel', 100)
+    else
+      vim.api.nvim_buf_set_lines(preview_bufnr, 0, -1, false, vim.fn.systemlist(string.format('bat %s', file_name)))
+    end
+  end,
+}
+
+
 previewers.Previewer = Previewer
 
 return previewers
diff --git a/lua/telescope/sorters.lua b/lua/telescope/sorters.lua
index a648007..c97a1a9 100644
--- a/lua/telescope/sorters.lua
+++ b/lua/telescope/sorters.lua
@@ -1,3 +1,5 @@
+local util = require('telescope.utils')
+
 local sorters = {}
 
 
@@ -31,4 +33,30 @@ end
 
 sorters.Sorter = Sorter
 
+sorters.get_ngram_sorter = function()
+  return Sorter:new {
+    scoring_function = function(_, prompt, line)
+      if prompt == "" or prompt == nil then
+        return 1
+      end
+
+      local ok, result = pcall(function()
+        local ngram = util.new_ngram { N = 4 }
+        ngram:add(line)
+
+        local score = ngram:score(prompt)
+        if score == 0 then
+          return -1
+        end
+
+        -- return math.pow(math.max(score, 0.0001), -1)
+        return score
+      end)
+
+      print(prompt, line, result)
+      return ok and result or 1
+    end
+  }
+end
+
 return sorters
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