From 54b25a873125aecd4e0e347858b39881361b3329 Mon Sep 17 00:00:00 2001
From: Luc Georges <McPatate@users.noreply.github.com>
Date: Wed, 7 Feb 2024 15:06:58 +0100
Subject: [PATCH] feat: add socket connection (#72)

---
 Cargo.lock                | 17 +++++++++--------
 crates/llm-ls/Cargo.toml  |  9 +++++----
 crates/llm-ls/src/main.rs | 34 +++++++++++++++++++++++++++++++---
 3 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 3129b6c..bd3f737 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -974,6 +974,7 @@ checksum = "da2479e8c062e40bf0066ffa0bc823de0a9368974af99c9f6df941d2c231e03f"
 name = "llm-ls"
 version = "0.4.0"
 dependencies = [
+ "clap",
  "home",
  "reqwest",
  "ropey",
@@ -2063,9 +2064,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
 
 [[package]]
 name = "tokenizers"
-version = "0.14.1"
+version = "0.15.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d9be88c795d8b9f9c4002b3a8f26a6d0876103a6f523b32ea3bac52d8560c17c"
+checksum = "6db445cceba5dfeb0f9702be7d6bfd91801ddcbe8fe8722defe7f2e96da75812"
 dependencies = [
  "aho-corasick",
  "derive_builder",
@@ -2368,9 +2369,9 @@ dependencies = [
 
 [[package]]
 name = "tree-sitter-erlang"
-version = "0.2.0"
+version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8d110d62a7ae35b985d8cfbc4de6e9281c7cbf268c466e30ebb31c2d3f861141"
+checksum = "93ced5145ebb17f83243bf055b74e108da7cc129e12faab4166df03f59b287f4"
 dependencies = [
  "cc",
  "tree-sitter",
@@ -2388,9 +2389,9 @@ dependencies = [
 
 [[package]]
 name = "tree-sitter-html"
-version = "0.19.0"
+version = "0.20.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "184e6b77953a354303dc87bf5fe36558c83569ce92606e7b382a0dc1b7443443"
+checksum = "017822b6bd42843c4bd67fabb834f61ce23254e866282dd93871350fd6b7fa1d"
 dependencies = [
  "cc",
  "tree-sitter",
@@ -2518,9 +2519,9 @@ dependencies = [
 
 [[package]]
 name = "tree-sitter-swift"
-version = "0.3.6"
+version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "eee2dbeb101a88a1d9e4883e3fbda6c799cf676f6a1cf59e4fc3862e67e70118"
+checksum = "452e6ee0a14b82a0dcd93400b8d3fe3784fdbd775191a89ef84586e5ccec6be7"
 dependencies = [
  "cc",
  "tree-sitter",
diff --git a/crates/llm-ls/Cargo.toml b/crates/llm-ls/Cargo.toml
index f70e218..dbc3324 100644
--- a/crates/llm-ls/Cargo.toml
+++ b/crates/llm-ls/Cargo.toml
@@ -7,6 +7,7 @@ edition = "2021"
 name = "llm-ls"
 
 [dependencies]
+clap = { version = "4", features = ["derive"] }
 home = "0.5"
 ropey = { version = "1.6", default-features = false, features = [
   "simd",
@@ -19,7 +20,7 @@ reqwest = { version = "0.11", default-features = false, features = [
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 thiserror = "1"
-tokenizers = { version = "0.14", default-features = false, features = ["onig"] }
+tokenizers = { version = "0.15", default-features = false, features = ["onig"] }
 tokio = { version = "1", features = [
   "fs",
   "io-std",
@@ -37,9 +38,9 @@ tree-sitter-c = "0.20"
 tree-sitter-cpp = "0.20"
 tree-sitter-c-sharp = "0.20"
 tree-sitter-elixir = "0.1"
-tree-sitter-erlang = "0.2"
+tree-sitter-erlang = "0.4"
 tree-sitter-go = "0.20"
-tree-sitter-html = "0.19"
+tree-sitter-html = "0.20"
 tree-sitter-java = "0.20"
 tree-sitter-javascript = "0.20"
 tree-sitter-json = "0.20"
@@ -52,7 +53,7 @@ tree-sitter-r = "0.19"
 tree-sitter-ruby = "0.20"
 tree-sitter-rust = "0.20"
 tree-sitter-scala = "0.20"
-tree-sitter-swift = "0.3"
+tree-sitter-swift = "0.4"
 tree-sitter-typescript = "0.20"
 
 [dependencies.uuid]
diff --git a/crates/llm-ls/src/main.rs b/crates/llm-ls/src/main.rs
index c6f2c54..cab0879 100644
--- a/crates/llm-ls/src/main.rs
+++ b/crates/llm-ls/src/main.rs
@@ -1,3 +1,4 @@
+use clap::Parser;
 use ropey::Rope;
 use serde::{Deserialize, Deserializer, Serialize};
 use serde_json::{Map, Value};
@@ -8,6 +9,7 @@ use std::sync::Arc;
 use std::time::{Duration, Instant};
 use tokenizers::Tokenizer;
 use tokio::io::AsyncWriteExt;
+use tokio::net::TcpListener;
 use tokio::sync::RwLock;
 use tower_lsp::jsonrpc::Result as LspResult;
 use tower_lsp::lsp_types::*;
@@ -770,10 +772,22 @@ impl LanguageServer for LlmService {
     }
 }
 
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+struct Args {
+    /// Wether to use a tcp socket for data transfer
+    #[arg(long = "port")]
+    socket: Option<usize>,
+
+    /// Wether to use stdio transport for data transfer, ignored because it is the default
+    /// behaviour
+    #[arg(short, long, default_value_t = true)]
+    stdio: bool,
+}
+
 #[tokio::main]
 async fn main() {
-    let stdin = tokio::io::stdin();
-    let stdout = tokio::io::stdout();
+    let args = Args::parse();
 
     let home_dir = home::home_dir().ok_or(()).expect("failed to find home dir");
     let cache_dir = home_dir.join(".cache/llm_ls");
@@ -822,5 +836,19 @@ async fn main() {
     .custom_method("llm-ls/rejectCompletion", LlmService::reject_completion)
     .finish();
 
-    Server::new(stdin, stdout, socket).serve(service).await;
+    if let Some(port) = args.socket {
+        let addr = format!("127.0.0.1:{port}");
+        let listener = TcpListener::bind(&addr)
+            .await
+            .unwrap_or_else(|_| panic!("failed to bind tcp listener to {addr}"));
+        let (stream, _) = listener
+            .accept()
+            .await
+            .unwrap_or_else(|_| panic!("failed to accept new connections on {addr}"));
+        let (read, write) = tokio::io::split(stream);
+        Server::new(read, write, socket).serve(service).await;
+    } else {
+        let (stdin, stdout) = (tokio::io::stdin(), tokio::io::stdout());
+        Server::new(stdin, stdout, socket).serve(service).await;
+    }
 }