<commit_before>// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use serde::Deserialize;
use serde::Serialize;
use serde_json::json;
use std::collections::HashMap;
use std::io;
use std::sync::mpsc::channel;
use std::sync::mpsc::Sender;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::Weak;

pub enum Message {
    Request(Request),
    Notification(Notification),
}

pub struct ResponseHandle {
    id: Id,
    writer: Arc<Mutex<Write + Send>>,
}

impl ResponseHandle {
    pub fn respond(self, response: Result<serde_json::Value, serde_json::Value>) {
        self.writer
            .lock()
            .unwrap()
            .write_packet(match response {
                Ok(result) => {
                    json!({ "jsonrpc": "2.0", "id": self.id, "result": result}).to_string()
                }
                Err(error) => json!({ "jsonrpc": "2.0", "id": self.id, "error": error}).to_string(),
            })
            .unwrap();
    }
}

pub struct Request {
    pub method: String,
    pub params: serde_json::Value,
    pub response_handle: ResponseHandle,
}

pub struct Notification {
    pub method: String,
    pub params: serde_json::Value,
}

type ResultOrError = Result<serde_json::Value, serde_json::Value>;
type MyMap = HashMap<Id, Sender<ResultOrError>>;

// Basic implementation of LSP Server.
//
// It exists after receiving `exit` notification and forwards all other requests and responses to
// handler passed in run method.
pub struct Server<R: Read, W: Write> {
    reader: R,
    writer: Arc<Mutex<W>>,
    // Map of requests that are currently waiting for the response from client.
    running_requests: Arc<Mutex<MyMap>>,
}

#[derive(Debug, PartialEq, Clone, Hash, Eq, Deserialize, Serialize)]
#[serde(untagged)]
enum Id {
    Number(i64),
    String(String),
}

impl<R, W> Iterator for Server<R, W>
where
    R: Read,
    W: Write + Send + 'static,
{
    type Item = Message;

    fn next(&mut self) -> Option<Message> {
        loop {
            let packet = match self.reader.read_packet() {
                Ok(packet) => packet,
                // TODO: Save the error
                Err(_) => return None,
            };
            let json: serde_json::Value = match serde_json::from_str(&packet) {
                Ok(value) => value,
                // TODO: We should probably reply with error?
                Err(_) => return None,
            };
            match &json {
                serde_json::Value::Object(map) => {
                    if let Some(serde_json::Value::String(method)) = map.get("method") {
                        if method == "exit" {
                            return None;
                        }
                        if let Some(id_val) = map.get("id") {
                            let id: Id = serde_json::from_value(id_val.clone()).unwrap();
                            return Some(Message::Request(Request {
                                method: method.to_string(),
                                params: json["params"].clone(),
                                response_handle: ResponseHandle {
                                    id: id,
                                    writer: self.writer.clone(),
                                },
                            }));
                        }
                        return Some(Message::Notification(Notification {
                            method: method.to_string(),
                            params: json["params"].clone(),
                        }));
                    }
                    if let Some(result) = map.get("result") {
                        if let Some(id_val) = map.get("id") {
                            let id: Id = serde_json::from_value(id_val.clone()).unwrap();
                            self.running_requests.lock().unwrap()[&id]
                                .send(Ok(result.clone()))
                                .unwrap();
                            continue;
                        }
                    }
                    // TODO: I think we should just respond with error here.
                    return None;
                }
                // TODO: I think we should just respond with error here.
                _ => return None,
            }
        }
    }
}

impl<R, W> Server<R, W>
where
    R: Read,
    W: Write + Send + 'static,
{
    pub fn new(reader: R, writer: W) -> Server<R, W> {
        return Server {
            reader: reader,
            writer: Arc::new(Mutex::new(writer)),
            running_requests: Arc::new(Mutex::new(HashMap::new())),
        };
    }

    pub fn sender(&self) -> MySender {
        return MySender {
            next_id: Arc::new(Mutex::new(Counter::new())),
            writer: self.writer.clone(),
            running_requests: Arc::downgrade(&self.running_requests),
        };
    }
}

struct Counter {
    id: i64,
}

impl Counter {
    fn new() -> Counter {
        Counter { id: 0 }
    }

    fn next(&mut self) -> i64 {
        self.id += 1;
        return self.id;
    }
}

pub struct MySender {
    next_id: Arc<Mutex<Counter>>,
    writer: Arc<Mutex<Write + Send>>,
    running_requests: Weak<Mutex<MyMap>>,
}

impl MySender {
    pub fn send_notification(&self, method: &str, params: serde_json::Value) {
        // TODO: how to properly handle errors here?
        self.writer
            .lock()
            .unwrap()
            .write_packet(
                json!(
                {
                    "jsonrpc": "2.0",
                    "method": method,
                    "params": params}
                )
                .to_string(),
            )
            .unwrap();<commit_msg>Add `dyn` to traits in server module.<commit_after>// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use serde::Deserialize;
use serde::Serialize;
use serde_json::json;
use std::collections::HashMap;
use std::io;
use std::sync::mpsc::channel;
use std::sync::mpsc::Sender;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::Weak;

pub enum Message {
    Request(Request),
    Notification(Notification),
}

pub struct ResponseHandle {
    id: Id,
    writer: Arc<Mutex<dyn Write + Send>>,
}

impl ResponseHandle {
    pub fn respond(self, response: Result<serde_json::Value, serde_json::Value>) {
        self.writer
            .lock()
            .unwrap()
            .write_packet(match response {
                Ok(result) => {
                    json!({ "jsonrpc": "2.0", "id": self.id, "result": result}).to_string()
                }
                Err(error) => json!({ "jsonrpc": "2.0", "id": self.id, "error": error}).to_string(),
            })
            .unwrap();
    }
}

pub struct Request {
    pub method: String,
    pub params: serde_json::Value,
    pub response_handle: ResponseHandle,
}

pub struct Notification {
    pub method: String,
    pub params: serde_json::Value,
}

type ResultOrError = Result<serde_json::Value, serde_json::Value>;
type MyMap = HashMap<Id, Sender<ResultOrError>>;

// Basic implementation of LSP Server.
//
// It exists after receiving `exit` notification and forwards all other requests and responses to
// handler passed in run method.
pub struct Server<R: Read, W: Write> {
    reader: R,
    writer: Arc<Mutex<W>>,
    // Map of requests that are currently waiting for the response from client.
    running_requests: Arc<Mutex<MyMap>>,
}

#[derive(Debug, PartialEq, Clone, Hash, Eq, Deserialize, Serialize)]
#[serde(untagged)]
enum Id {
    Number(i64),
    String(String),
}

impl<R, W> Iterator for Server<R, W>
where
    R: Read,
    W: Write + Send + 'static,
{
    type Item = Message;

    fn next(&mut self) -> Option<Message> {
        loop {
            let packet = match self.reader.read_packet() {
                Ok(packet) => packet,
                // TODO: Save the error
                Err(_) => return None,
            };
            let json: serde_json::Value = match serde_json::from_str(&packet) {
                Ok(value) => value,
                // TODO: We should probably reply with error?
                Err(_) => return None,
            };
            match &json {
                serde_json::Value::Object(map) => {
                    if let Some(serde_json::Value::String(method)) = map.get("method") {
                        if method == "exit" {
                            return None;
                        }
                        if let Some(id_val) = map.get("id") {
                            let id: Id = serde_json::from_value(id_val.clone()).unwrap();
                            return Some(Message::Request(Request {
                                method: method.to_string(),
                                params: json["params"].clone(),
                                response_handle: ResponseHandle {
                                    id: id,
                                    writer: self.writer.clone(),
                                },
                            }));
                        }
                        return Some(Message::Notification(Notification {
                            method: method.to_string(),
                            params: json["params"].clone(),
                        }));
                    }
                    if let Some(result) = map.get("result") {
                        if let Some(id_val) = map.get("id") {
                            let id: Id = serde_json::from_value(id_val.clone()).unwrap();
                            self.running_requests.lock().unwrap()[&id]
                                .send(Ok(result.clone()))
                                .unwrap();
                            continue;
                        }
                    }
                    // TODO: I think we should just respond with error here.
                    return None;
                }
                // TODO: I think we should just respond with error here.
                _ => return None,
            }
        }
    }
}

impl<R, W> Server<R, W>
where
    R: Read,
    W: Write + Send + 'static,
{
    pub fn new(reader: R, writer: W) -> Server<R, W> {
        return Server {
            reader: reader,
            writer: Arc::new(Mutex::new(writer)),
            running_requests: Arc::new(Mutex::new(HashMap::new())),
        };
    }

    pub fn sender(&self) -> MySender {
        return MySender {
            next_id: Arc::new(Mutex::new(Counter::new())),
            writer: self.writer.clone(),
            running_requests: Arc::downgrade(&self.running_requests),
        };
    }
}

struct Counter {
    id: i64,
}

impl Counter {
    fn new() -> Counter {
        Counter { id: 0 }
    }

    fn next(&mut self) -> i64 {
        self.id += 1;
        return self.id;
    }
}

pub struct MySender {
    next_id: Arc<Mutex<Counter>>,
    writer: Arc<Mutex<dyn Write + Send>>,
    running_requests: Weak<Mutex<MyMap>>,
}

impl MySender {
    pub fn send_notification(&self, method: &str, params: serde_json::Value) {
        // TODO: how to properly handle errors here?
        self.writer
            .lock()
            .unwrap()
            .write_packet(
                json!(
                {
                    "jsonrpc": "2.0",
                    "method": method,
                    "params": params}
                )
                .to_string(),
            )
            .unwrap();
{'fraction_non_alphanumeric': 0.12176498955151904, 'fraction_numerical': 0.00353640893746986, 'mean_word_length': 1.5687448389760528, 'pattern_counts': {'":': 18, '<': 49, '<?xml version=': 0, '>': 75, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2462328', 'n_tokens_mistral': 3413, 'n_tokens_neox': 3340, 'n_words': 931}