<commit_before>﻿using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using System.Threading.Tasks.Dataflow;
using Microsoft.Common.Core;
using Microsoft.Languages.Editor.Shell;
using Microsoft.R.Host.Client;
using Microsoft.R.Support.Settings;
using Microsoft.VisualStudio.Shell;
using Task = System.Threading.Tasks.Task;
using Microsoft.VisualStudio.R.Package.Plots;
using Microsoft.VisualStudio.Shell.Interop;
using Microsoft.VisualStudio.R.Package.Shell;

namespace Microsoft.VisualStudio.R.Package.Repl.Session {
    internal sealed class RSession : IRSession, IRCallbacks {
        private readonly BufferBlock<RSessionRequestSource> _pendingRequestSources = new BufferBlock<RSessionRequestSource>();
        private readonly ConcurrentQueue<RSessionEvaluationSource> _pendingEvaluationSources = new ConcurrentQueue<RSessionEvaluationSource>();
        private readonly Stack<RSessionRequestSource> _currentRequestSources = new Stack<RSessionRequestSource>();

        public event EventHandler<RBeforeRequestEventArgs> BeforeRequest;
        public event EventHandler<RResponseEventArgs> Response;
        public event EventHandler<RErrorEventArgs> Error;
        public event EventHandler<EventArgs> Disconnected;

        /// <summary>
        /// ReadConsole requires a task even if there are no pending requests
        /// </summary>
        private IReadOnlyCollection<IRContext> _contexts;
        private RHost _host;
        private Task _hostRunTask;
        private TaskCompletionSource<object> _initializationTcs;

        public string Prompt { get; private set; } = "> ";
        public int MaxLength { get; private set; } = 0x1000;
        public bool IsHostRunning => _hostRunTask != null && !_hostRunTask.IsCompleted;

        public void Dispose() {
            _host?.Dispose();
        }

        public Task<IRSessionInteraction> BeginInteractionAsync(bool isVisible = true) {
            RSessionRequestSource requestSource = new RSessionRequestSource(isVisible, _contexts);
            _pendingRequestSources.Post(requestSource);
            return requestSource.CreateRequestTask;
        }

        public Task<IRSessionEvaluation> BeginEvaluationAsync() {
            var source = new RSessionEvaluationSource();
            _pendingEvaluationSources.Enqueue(source);
            return source.Task;
        }

        public Task StartHostAsync() {
            if (_hostRunTask != null && !_hostRunTask.IsCompleted) {
                throw new InvalidOperationException("Another instance of RHost is running for this RSession. Stop it before starting new one.");
            }

            _host = new RHost(this);
            _initializationTcs = new TaskCompletionSource<object>();

            _hostRunTask = Task.Run(() => _host.CreateAndRun(RToolsSettings.GetRVersionPath()));

            var initializationTask = _initializationTcs.Task.ContinueWith(new Func<Task, Task>(AfterInitialization)).Unwrap();

            return Task.WhenAny(initializationTask, _hostRunTask).Unwrap();
        }

        public async Task StopHostAsync() {
            if (_hostRunTask.IsCompleted) {
                return;
            }

            var request = await BeginInteractionAsync(false);
            if (_hostRunTask.IsCompleted) {
                request.Dispose();
                return;
            }

            await request.Quit();
            await _hostRunTask;
        }

        private async Task AfterInitialization(Task task) {
            var interaction = await BeginInteractionAsync(false);
            await interaction.SetDefaultWorkingDirectory();
        }

        Task IRCallbacks.Connected(string rVersion) {
            _initializationTcs.SetResult(null);
            return Task.CompletedTask;
        }

        Task IRCallbacks.Disconnected() {
            while (_currentRequestSources.Count > 0) {
                var requestSource = _currentRequestSources.Pop();
                requestSource.Complete();
            }
            while (_pendingRequestSources.Count > 0) {
                var requestSource = _pendingRequestSources.Receive();
                requestSource.Complete();
            }
            while (_pendingEvaluationSources.Count > 0) {
                RSessionEvaluationSource evalSource;
                if (_pendingEvaluationSources.TryDequeue(out evalSource)) {
                    evalSource.TryCancel();
                }
            }

            _contexts = null;

            OnDisconnected();
            return Task.CompletedTask;
        }

        async Task<string> IRCallbacks.ReadConsole(IReadOnlyCollection<IRContext> contexts, string prompt, string buf, int len, bool addToHistory, CancellationToken ct) {
            foreach (var rsToCompleter in _currentRequestSources.PopWhile(rs => rs.Contexts.Count >= contexts.Count)) {
                rsToCompleter.Complete();
            }

            _contexts = contexts;
            Prompt = prompt;
            MaxLength = len;

            OnBeforeRequest(contexts, prompt, len, addToHistory);

            while (true) {
                ct.ThrowIfCancellationRequested();

                try {
                    return await ReadNextRequest(prompt, len, ct);
                } catch (TaskCanceledException) {
                    // If request was canceled through means other than our token, it indicates the refusal of
                    // that requestor to respond to that particular prompt, so move on to the next requestor.
                    // If it was canceled through the token, then host itself is shutting down, and cancellation
                    // will be propagated on the entry to next iteration of this loop.
                }
            }
        }

        private async Task<string> ReadNextRequest(string prompt, int len, CancellationToken ct) {
            var requestSource = await _pendingRequestSources.ReceiveAsync(ct);

            TaskCompletionSource<string> requestTcs = new TaskCompletionSource<string>();
            _currentRequestSources.Push(requestSource);
            requestSource.Request(prompt, len, requestTcs);

            // Cancel immediately if ct is canceled, even if request is not handled yet.
            // Disconnected handler will run next and take care of cleaning the queue.
            await Task.WhenAny(requestTcs.Task, Task.Delay(Timeout.Infinite, ct)).Unwrap();
            string response = await requestTcs.Task;

            Debug.Assert(response.Length < len); // len includes null terminator
            if (response.Length >= len) {
                response = response.Substring(0, len - 1);
            }

            return response;
        }

        Task IRCallbacks.WriteConsoleEx(IReadOnlyCollection<IRContext> contexts, string buf, OutputType otype, CancellationToken ct) {
            if (otype == OutputType.Error) {<commit_msg>Simplify ReadNextRequest.<commit_after>﻿using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using System.Threading.Tasks.Dataflow;
using Microsoft.Common.Core;
using Microsoft.Languages.Editor.Shell;
using Microsoft.R.Host.Client;
using Microsoft.R.Support.Settings;
using Microsoft.VisualStudio.Shell;
using Task = System.Threading.Tasks.Task;
using Microsoft.VisualStudio.R.Package.Plots;
using Microsoft.VisualStudio.Shell.Interop;
using Microsoft.VisualStudio.R.Package.Shell;

namespace Microsoft.VisualStudio.R.Package.Repl.Session {
    internal sealed class RSession : IRSession, IRCallbacks {
        private readonly BufferBlock<RSessionRequestSource> _pendingRequestSources = new BufferBlock<RSessionRequestSource>();
        private readonly ConcurrentQueue<RSessionEvaluationSource> _pendingEvaluationSources = new ConcurrentQueue<RSessionEvaluationSource>();
        private readonly Stack<RSessionRequestSource> _currentRequestSources = new Stack<RSessionRequestSource>();

        public event EventHandler<RBeforeRequestEventArgs> BeforeRequest;
        public event EventHandler<RResponseEventArgs> Response;
        public event EventHandler<RErrorEventArgs> Error;
        public event EventHandler<EventArgs> Disconnected;

        /// <summary>
        /// ReadConsole requires a task even if there are no pending requests
        /// </summary>
        private IReadOnlyCollection<IRContext> _contexts;
        private RHost _host;
        private Task _hostRunTask;
        private TaskCompletionSource<object> _initializationTcs;

        public string Prompt { get; private set; } = "> ";
        public int MaxLength { get; private set; } = 0x1000;
        public bool IsHostRunning => _hostRunTask != null && !_hostRunTask.IsCompleted;

        public void Dispose() {
            _host?.Dispose();
        }

        public Task<IRSessionInteraction> BeginInteractionAsync(bool isVisible = true) {
            RSessionRequestSource requestSource = new RSessionRequestSource(isVisible, _contexts);
            _pendingRequestSources.Post(requestSource);
            return requestSource.CreateRequestTask;
        }

        public Task<IRSessionEvaluation> BeginEvaluationAsync() {
            var source = new RSessionEvaluationSource();
            _pendingEvaluationSources.Enqueue(source);
            return source.Task;
        }

        public Task StartHostAsync() {
            if (_hostRunTask != null && !_hostRunTask.IsCompleted) {
                throw new InvalidOperationException("Another instance of RHost is running for this RSession. Stop it before starting new one.");
            }

            _host = new RHost(this);
            _initializationTcs = new TaskCompletionSource<object>();

            _hostRunTask = Task.Run(() => _host.CreateAndRun(RToolsSettings.GetRVersionPath()));

            var initializationTask = _initializationTcs.Task.ContinueWith(new Func<Task, Task>(AfterInitialization)).Unwrap();

            return Task.WhenAny(initializationTask, _hostRunTask).Unwrap();
        }

        public async Task StopHostAsync() {
            if (_hostRunTask.IsCompleted) {
                return;
            }

            var request = await BeginInteractionAsync(false);
            if (_hostRunTask.IsCompleted) {
                request.Dispose();
                return;
            }

            await request.Quit();
            await _hostRunTask;
        }

        private async Task AfterInitialization(Task task) {
            var interaction = await BeginInteractionAsync(false);
            await interaction.SetDefaultWorkingDirectory();
        }

        Task IRCallbacks.Connected(string rVersion) {
            _initializationTcs.SetResult(null);
            return Task.CompletedTask;
        }

        Task IRCallbacks.Disconnected() {
            while (_currentRequestSources.Count > 0) {
                var requestSource = _currentRequestSources.Pop();
                requestSource.Complete();
            }
            while (_pendingRequestSources.Count > 0) {
                var requestSource = _pendingRequestSources.Receive();
                requestSource.Complete();
            }
            while (_pendingEvaluationSources.Count > 0) {
                RSessionEvaluationSource evalSource;
                if (_pendingEvaluationSources.TryDequeue(out evalSource)) {
                    evalSource.TryCancel();
                }
            }

            _contexts = null;

            OnDisconnected();
            return Task.CompletedTask;
        }

        async Task<string> IRCallbacks.ReadConsole(IReadOnlyCollection<IRContext> contexts, string prompt, string buf, int len, bool addToHistory, CancellationToken ct) {
            foreach (var rsToCompleter in _currentRequestSources.PopWhile(rs => rs.Contexts.Count >= contexts.Count)) {
                rsToCompleter.Complete();
            }

            _contexts = contexts;
            Prompt = prompt;
            MaxLength = len;

            OnBeforeRequest(contexts, prompt, len, addToHistory);

            while (true) {
                ct.ThrowIfCancellationRequested();

                try {
                    return await ReadNextRequest(prompt, len, ct);
                } catch (TaskCanceledException) {
                    // If request was canceled through means other than our token, it indicates the refusal of
                    // that requestor to respond to that particular prompt, so move on to the next requestor.
                    // If it was canceled through the token, then host itself is shutting down, and cancellation
                    // will be propagated on the entry to next iteration of this loop.
                }
            }
        }

        private async Task<string> ReadNextRequest(string prompt, int len, CancellationToken ct) {
            var requestSource = await _pendingRequestSources.ReceiveAsync(ct);

            TaskCompletionSource<string> requestTcs = new TaskCompletionSource<string>();
            _currentRequestSources.Push(requestSource);
            requestSource.Request(prompt, len, requestTcs);
            ct.Register(delegate { requestTcs.TrySetCanceled(); });

            string response = await requestTcs.Task;

            Debug.Assert(response.Length < len); // len includes null terminator
            if (response.Length >= len) {
                response = response.Substring(0, len - 1);
            }

            return response;
        }

        Task IRCallbacks.WriteConsoleEx(IReadOnlyCollection<IRContext> contexts, string buf, OutputType otype, CancellationToken ct) {
            if (otype == OutputType.Error) {
{'fraction_non_alphanumeric': 0.07823819182845553, 'fraction_numerical': 0.0014488554042306578, 'mean_word_length': 2.394393902139169, 'pattern_counts': {'":': 0, '<': 53, '<?xml version=': 0, '>': 69, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3300983', 'n_tokens_mistral': 3388, 'n_tokens_neox': 3219, 'n_words': 846}