From 8186e6025a368f476b8562cf1c313e973dbf1971 Mon Sep 17 00:00:00 2001
From: Aaron Pham <contact@aarnphm.xyz>
Date: Thu, 15 Aug 2024 08:39:42 -0400
Subject: [PATCH] feat(tiktoken): automatic build

Signed-off-by: Aaron Pham <contact@aarnphm.xyz>
---
 Makefile             | 52 +++++++++++++++++++++++++++++++++++++++++++-
 README.md            | 30 ++++++++-----------------
 lua/avante/init.lua  |  1 +
 lua/tiktoken_lib.lua | 29 ++++++++++++++++++++++++
 4 files changed, 90 insertions(+), 22 deletions(-)
 create mode 100644 lua/tiktoken_lib.lua

diff --git a/Makefile b/Makefile
index 6ff648f55..d12f20cf4 100644
--- a/Makefile
+++ b/Makefile
@@ -1,5 +1,55 @@
+UNAME := $(shell uname)
+ARCH := $(shell uname -m)
+
+ifeq ($(UNAME), Linux)
+	OS := linux
+	EXT := so
+else ifeq ($(UNAME), Darwin)
+	OS := macOS
+	EXT := dylib
+else
+	$(error Unsupported operating system: $(UNAME))
+endif
+
+LUA_VERSIONS := luajit lua51
+BUILD_DIR := build
+
+all: luajit
+
+luajit: $(BUILD_DIR)/tiktoken_core.$(EXT)
+lua51: $(BUILD_DIR)/tiktoken_core-lua51.$(EXT)
+
+define build_from_source
+	git clone https://github.com/gptlang/lua-tiktoken.git $(BUILD_DIR)/lua-tiktoken-temp
+	cd $(BUILD_DIR)/lua-tiktoken-temp && cargo build --features=$1
+	cp $(BUILD_DIR)/lua-tiktoken-temp/target/debug/libtiktoken_core.$(EXT) $(BUILD_DIR)/tiktoken_core.$(EXT)
+	rm -rf $(BUILD_DIR)/lua-tiktoken-temp
+endef
+
+define download_release
+	curl -L https://github.com/gptlang/lua-tiktoken/releases/latest/download/tiktoken_core-$1-$2.$(EXT) -o $(BUILD_DIR)/tiktoken_core.$(EXT)
+endef
+
+ifeq ($(ARCH), arm64)
+    $(BUILD_DIR)/tiktoken_core.$(EXT): $(BUILD_DIR)
+	$(call build_from_source,luajit)
+    $(BUILD_DIR)/tiktoken_core-lua51.$(EXT): $(BUILD_DIR)
+	$(call build_from_source,lua51)
+else
+    $(BUILD_DIR)/tiktoken_core.$(EXT): $(BUILD_DIR)
+	$(call download_release,$(OS),luajit)
+    $(BUILD_DIR)/tiktoken_core-lua51.$(EXT): $(BUILD_DIR)
+	$(call download_release,$(OS),lua51)
+endif
+
+$(BUILD_DIR):
+	mkdir -p $(BUILD_DIR)
+
+clean:
+	rm -rf $(BUILD_DIR)
+
 luacheck:
-	luacheck `find -name  "*.lua"` --codes
+	luacheck `find -name "*.lua"` --codes
 
 stylecheck:
 	stylua --check lua/
diff --git a/README.md b/README.md
index a6c9cdadf..11554cbf0 100644
--- a/README.md
+++ b/README.md
@@ -20,31 +20,14 @@ https://github.com/user-attachments/assets/86140bfd-08b4-483d-a887-1b701d9e37dd
 
 ## Installation
 
-1. Install `tiktoken_core` (Optional):
-
-    - tiktoken_core: `sudo luarocks install --lua-version 5.1 tiktoken_core`. Alternatively, download a pre-built binary from [lua-tiktoken releases](https://github.com/gptlang/lua-tiktoken/releases)
-
-    - You can check your Lua PATH in Neovim by doing `:lua print(package.cpath)`. Save the binary as `tiktoken_core.so` in any of the given paths.
-
-    > For Arch Linux user, you can install [`luajit-tiktoken-bin`](https://aur.archlinux.org/packages/luajit-tiktoken-bin) or [`lua51-tiktoken-bin`](https://aur.archlinux.org/packages/lua51-tiktoken-bin) from aur!
-    >
-    > For macOS users:
-    > 1. Install `luarocks` using Homebrew: `brew install luarocks`
-    > 2. Install `tiktoken_core`: `luarocks install --lua-version 5.1 tiktoken_core`
-    > 3. Copy the installed `tiktoken_core.so` to your Homebrew Lua path:
-    >    `cp ~/.luarocks/lib/lua/5.1/tiktoken_core.so $(brew --prefix)/lib/lua/5.1/`
-    >
-    > This ensures `tiktoken_core` is properly installed and accessible in your Neovim environment.
-
-2. Install `avante.nvim` using [lazy.nvim](https://github.com/folke/lazy.nvim):
+Install `avante.nvim` using [lazy.nvim](https://github.com/folke/lazy.nvim):
 
     ```lua
     {
       "yetone/avante.nvim",
       event = "VeryLazy",
-      config = function()
-          require("avante").setup({})
-      end,
+      opts = {},
+      build = "make",
       dependencies = {
         "nvim-tree/nvim-web-devicons",
         {
@@ -59,7 +42,12 @@ https://github.com/user-attachments/assets/86140bfd-08b4-483d-a887-1b701d9e37dd
     }
     ```
 
-3. Default setup configuration:
+> [!IMPORTANT]
+>
+> If your neovim doesn't use LuaJIT, then change `build` to `make lua51`. By default running make will install luajit.
+> For ARM-based setup, make sure to also install cargo as we will have to build the tiktoken_core from source.
+
+Default setup configuration:
 
     ```lua
     {
diff --git a/lua/avante/init.lua b/lua/avante/init.lua
index 667297c05..98e1eb351 100644
--- a/lua/avante/init.lua
+++ b/lua/avante/init.lua
@@ -3,6 +3,7 @@ local sidebar = require("avante.sidebar")
 local config = require("avante.config")
 
 function M.setup(opts)
+  require("tiktoken_lib").load()
   config.update(opts)
   sidebar.setup()
 end
diff --git a/lua/tiktoken_lib.lua b/lua/tiktoken_lib.lua
new file mode 100644
index 000000000..f2811b666
--- /dev/null
+++ b/lua/tiktoken_lib.lua
@@ -0,0 +1,29 @@
+local H = {}
+local M = {}
+
+H.get_os_name = function()
+  local os_name = vim.loop.os_uname().sysname
+  if os_name == "Linux" then
+    return "linux"
+  elseif os_name == "Darwin" then
+    return "macOS"
+  else
+    error("Unsupported operating system: " .. os_name)
+  end
+end
+
+H.library_path = function()
+  local os_name = H.get_os_name()
+  local ext = os_name == "linux" and "so" or "dylib"
+  local dirname = string.sub(debug.getinfo(1).source, 2, #"/tiktoken_lib.lua" * -1)
+  return dirname .. ("../build/?.%s"):format(ext)
+end
+
+M.load = function()
+  local library_path = H.library_path()
+  if not string.find(package.cpath, library_path, 1, true) then
+    package.cpath = package.cpath .. ";" .. library_path
+  end
+end
+
+return M