tsql.nvim/fennel/tsql.fnl

;; basically transpiled from Fennel's website https://fennel-lang.org/see
;; Lua (git:603e6cc, ->fennel->lua to remove docs): (reduced) 4000 GPT-3 tokens, 3800 Codex tokens, 9800 chars
;; Fennel (git:603e6cc): 5106 GPT-3 tokens, 3488 Codex tokens, 10110 chars

;; Fennel got hit mostly by white-spaces in GPT-3
;; Fennel (git:603e6cc, manually removed whitespace): 4048 GPT-3 tokens, 3443 Codex tokens, 9051 chars

(set buf-select.QBuf {})
(set buf-select.QBuf.__index buf-select.QBuf)
(fn buf-select.QBuf.new [self bufnr path filetype lang is-loaded]
  (assert (= (type bufnr) :number) "bufnr must be a number")
  (assert (= (type path) :string) "path must be a string")
  (assert (= (type filetype) :string) "filetype must be a string")
  (local qbuf {: bufnr : filetype :is_loaded is-loaded : lang : path})
  (setmetatable qbuf self)
  (set qbuf.__index self)
  qbuf)
(fn get-lang [bufnr]
  (let [(status lang) (pcall (fn [] (: (vim.treesitter.get_parser bufnr) :lang)))]
    (when (not status) (local path (vim.api.nvim_buf_get_name bufnr))
      (local err lang)
      (error (string.format "Error determining language for buffer %d: %s\n%s"
        bufnr path err)))
    lang))
(fn buf-select.QBuf.from_nvim_bufnr [bufnr]
  (let [path (vim.api.nvim_buf_get_name bufnr)
        filetype (vim.api.nvim_buf_get_option bufnr :filetype)]
    (assert (= (type filetype) :string))
    (when (= (length filetype) 0) (lua "return nil"))
    (local (sts _lang) (pcall get-lang bufnr))
    (var lang nil)
    (when sts (set lang _lang))
    (local is-loaded (vim.api.nvim_buf_is_loaded bufnr))
    (buf-select.QBuf:new bufnr path filetype lang is-loaded)))
(set buf-select.buf_match {})
(set buf-select.BufMatch {})
(set buf-select.BufMatch.__index buf-select.BufMatch)
(fn buf-select.buf_match.is_loaded []
  (buf-select.BufMatch.new (fn [buf] buf.is_loaded)))
(fn buf-select.BufMatch.new [match-fn]
  (let [self (setmetatable {} buf-select.BufMatch)]
    (set self.__index buf-select.BufMatch)
    (set self.match_fn match-fn)
    self))
(fn buf-select.buf_match.filetype [...]
  (let [filetypes [...]]
    (buf-select.BufMatch.new (fn [buf]
       (each [_ filetype (ipairs filetypes)]
         (when (= buf.filetype filetype)
           (lua "return true")))
       false))))
(fn buf-select.buf_match.any []
  (buf-select.BufMatch.new (fn [_] true)))
(fn buf-select.buf_match.path [...]
  (let [paths [...]]
    (buf-select.BufMatch.new (fn [buf]
       (each [_ path (ipairs paths)]
         (when (not= (string.find buf.path path) nil)
           (lua "return true")))
       false))))
(fn buf-select.buf_match.path_or [...] (buf-select.buf_match.path ...))
(fn buf-select.buf_match.path_and [...]
  (let [paths [...]]
    (buf-select.BufMatch.new (fn [buf]
       (each [_ path (ipairs paths)]
         (when (= (string.find buf.path path) nil)
           (lua "return false")))
       true))))
(fn buf-select.buf_match.ext [...]
  (let [exts [...]]
    (buf-select.BufMatch.new (fn [buf]
       (each [_ ext (ipairs exts)]
         (when (= (buf.path:sub (- (length ext))) ext)
           (lua "return true")))
       false))))
(fn buf-select.BufMatch.or_ [self q]
  (buf-select.BufMatch.new (fn [buf]
       (or (self.matched_fn buf) (q.matched_fn buf)))))
(fn buf-select.BufMatch.then_ [self q]
  (buf-select.BufMatch.new (fn [buf]
       (and (self.matched_fn buf) (q.matched_fn buf)))))
(fn buf-select.BufMatch.not_ [self]
  (buf-select.BufMatch.new (fn [buf] (not (self.matched_fn buf)))))
(fn buf-select.BufMatch.filter_on [self itr]
  (let [matched {}]
    (each [_ buf (ipairs itr)]
      (when (self.match_fn buf) (table.insert matched buf)))
    matched))
(fn list-bufs [] (vim.api.nvim_list_bufs))
(fn buf-select.nvim_get_qbufs []
  (let [bufnrs (list-bufs)
        qbufs {}]
    (each [_ bufnr (ipairs bufnrs)]
      (local qbuf (buf-select.QBuf.from_nvim_bufnr bufnr))
      (when (not= qbuf nil) (table.insert qbufs qbuf)))
    qbufs))
(local token-select {})
(set token-select.ts_query {})
(set token-select.TSQuery {})
(set token-select.TSQuery.__index token-select.TSQuery)
(fn token-select.TSQuery.new [q]
  (let [self q] (setmetatable self token-select.TSQuery)
    (set self.__index token-select.TSQuery)
    self))
(set token-select.FileLoc {})
(set token-select.FileLoc.__index token-select.FileLoc)
(fn token-select.FileLoc.new [file-loc] (assert (not= file-loc nil))
  (local self file-loc)
  (setmetatable self token-select.FileLoc)
  (set self.__index token-select.FileLoc)
  self)
(set token-select.QNode {})
(set token-select.QNode.__index token-select.QNode)
(fn token-select.QNode.new [qnode] (assert (not= qnode nil)) (local self qnode)
  (setmetatable self token-select.QNode)
  (set self.__index token-select.QNode)
  self)
(fn token-select.ts_query.from_scm [treesitter-query]
  (token-select.TSQuery.new {:query treesitter-query}))
(fn token-select.TSQuery.find_nodes [self files]
  (let [result {}]
    (each [_ file (ipairs files)]
      (when (not= file.lang nil)
        (local (sts parser) (pcall vim.treesitter.get_parser file.bufnr
           file.filetype))
        (when sts
          (local tree (. (parser:parse) 1))
          (local root (tree:root))
          (local (sts query) (pcall vim.treesitter.query.parse file.lang
            self.query))
          (when (not sts)
            (error (.. "Error parsing query \"" self.query "\" on lang "
               file.lang " for file " file.path)))
          (each [_ ___match___ _ (query:iter_matches root file.bufnr 0 (- 1))]
            (each [id node (pairs ___match___)]
              (local (start-row start-col end-row end-col) (node:range false))
              (local start
                 (token-select.FileLoc.new {:col_0 start-col
                  :row_0 start-row}))
              (local end-ex-col
                 (token-select.FileLoc.new {:col_0 end-col :row_0 end-row}))
              (local qnode
                 (token-select.QNode.new {:buf file
                  :end_ex_col end-ex-col
                  : start}))
              (table.insert result qnode))))))
    result))
(local M {})
(set M.sink_by {})
(set M.format {})
(set M.Sink {})
(set M.Sink.__index M.Sink)
(fn M.sink_by.pure_fn [func]
  (let [self (setmetatable {:sink (fn [_ nodes] (func nodes))} M.Sink)]
    (set self.__index M.Sink)
    self))
(fn M.sink_by.highlight []
  (M.sink_by.pure_fn (fn [nodes]
     (each [_ node (ipairs nodes)]
       (vim.highlight.range node.buf.bufnr M.config.nvim_ns
          M.config.nvim_hl_group
          [node.start.row_0
           node.start.col_0]
          [node.end_ex_col.row_0
           node.end_ex_col.col_0]
          {:inclusive false})))))
(fn M.format.display [nodes]
  (let [texts {}]
    (each [_ node (ipairs nodes)]
      (local text
         (vim.api.nvim_buf_get_text node.buf.bufnr node.start.row_0
            node.start.col_0 node.end_ex_col.row_0
            node.end_ex_col.col_0 {}))
      (table.insert texts (table.concat text "\n")))
    (table.concat texts "\n\n")))
(fn M.format.dump [nodes] (vim.inspect nodes {:indent "  " :newline "\n"}))
(set M.format.debug M.format.dump)
(set M.format.default M.format.display)
(fn M.sink_by.print [format]
  (when (= format nil) (set-forcibly! format M.format.default))
  (M.sink_by.pure_fn (fn [nodes] (print (format nodes)))))
(fn M.sink_by.nvim_yank_buf [format]
  (when (= format nil) (set-forcibly! format M.format.default))
  (M.sink.pure_fn (fn [nodes]
      (let [text (format nodes)] (vim.fn.setreg "\"" text)))))
(set M.buf_match buf-select.buf_match)
(set M.BufMatch buf-select.BufMatch)
(set M.QBuf buf-select.QBuf)
(set M.nvim_get_qbufs buf-select.nvim_get_qbufs)
(set M.ts_query token-select.ts_query)
(set M.TSQuery token-select.TSQuery)
(set M.Tsql {})
(set M.Tsql.__index M.Tsql)
(fn M.s [external-dsl])
(fn M.t [buf-match codeql sink]
  (setmetatable {:buf_match buf-match : codeql : sink} M.Tsql))
(fn M.Tsql.q_nodes [self] (self.codeql:find_nodes (self:qbufs)))
(fn M.Tsql.qbufs [self]
  (self.buf_match.filter_on self.buf_match (M.nvim_get_qbufs)))
(fn M.Tsql.do_nvim [self store] (self.sink:sink (self:q_nodes))
  (store:add_highlight self))
(fn M._delete_all_highlights [config store]
  (let [bufs {}]
    (each [_ highlight-q (pairs store.highlighting)]
      (each [_ qnode (ipairs (highlight-q:q_nodes))]
        (table.insert bufs qnode.buf.bufnr qnode.buf)))
    (each [_ buf (pairs bufs)]
      (vim.api.nvim_buf_clear_namespace buf.bufnr config.nvim_ns 0 (- 1)))
    (store:clear_highlights)))
(local Config {})
(set M.config_default {:nvim_hl_group :Search})
(set M.RtConfig {})
(set M.config {})
(set M.Store {})
(set M.Store.__index M.Store)
(fn M.Store.add_highlight [self tsql] (table.insert self.highlighting tsql))
(fn M.Store.clear_highlights [self] (set self.highlighting {}))
(fn M.Store.new [self]
  (let [o {:highlighting {}}]
    (setmetatable o self)
    (set o.__index self)
    o))
(fn M.clear_highlights [] (M._delete_all_highlights M.config M.store))
(fn M.setup [config]
  (set M.config
     (vim.tbl_deep_extend :force M.config (or config M.config_default)))
  (set M.config.nvim_ns (vim.api.nvim_create_namespace :tsql))
  (vim.api.nvim_create_user_command :Noh (fn [_] (M.clear_highlights))
    {:desc "Clear tsql highlights" :nargs 0})
  (vim.api.nvim_create_user_command :Tdsl
    (fn [cmd]
      (: (M.s cmd.args) :do_nvim M.store))
    {:desc "tsql DSL invocation" :nargs "?"})
  (set M.store (M.Store:new)))
M