LINQ Query Window for Unity

QueryWindow.cs

using Microsoft.CSharp;
using System;
using System.CodeDom.Compiler;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
using System.Text;
using System.Text.RegularExpressions;
using UnityEditor;
using UnityEngine;
using UnityEngine.Experimental.UIElements;
using UnityEngine.Experimental.UIElements.StyleEnums;
using UnityEngine.SceneManagement;

public sealed class QueryWindow : EditorWindow
{
    private static string[] ignoredNamespaces = { "System.Runtime.DesignerServices" };

    private static PropertyInfo[] ignoredProperties = { typeof(Renderer).GetProperty("material"),
        typeof(Renderer).GetProperty("materials") };

    private List<string> autoComplete = new List<string>();
    private int filterLength;
    private TextField queryField;
    private List<object> queryResult;

    [MenuItem("Window/Query")]
    public static void Get()
    {
        GetWindow<QueryWindow>("Query");
    }

    private void OnEnable()
    {
        queryField = new TextField();
        queryField.multiline = true;
        queryField.style.flexGrow = 1f;
        queryField.style.font = Font.CreateDynamicFontFromOSFont("Courier New", 13);
        queryField.value = "return Selection.objects;";
        queryField.RegisterCallback(new EventCallback<InputEvent>(queryField_Input));
        queryField.RegisterCallback(new EventCallback<KeyDownEvent>(queryField_KeyDown));

        var runButton = new Button(runButton_Click);
        runButton.text = "Run\n(F5)";

        var topRow = new VisualElement();
        topRow.style.flexDirection = FlexDirection.Row;
        topRow.Add(queryField);
        topRow.Add(runButton);

        var resultContainer = new IMGUIContainer(resultContainer_OnGUI);
        resultContainer.style.flexGrow = 1;

        rootVisualElement.Add(topRow);
        rootVisualElement.Add(resultContainer);
    }

    private void queryField_Input(InputEvent e)
    {
        if (e.newData.Length > e.previousData.Length)
        {
            int cursorIndex = queryField.cursorIndex;
            if (cursorIndex > 0 && e.newData[cursorIndex - 1] == '\n')
            {
                int startOfPrevLine = cursorIndex > 1
                        ? e.newData.LastIndexOf('\n', cursorIndex - 2) + 1 : 0;

                int indent = startOfPrevLine;
                while (e.newData[indent] == ' ')
                    indent++;
                indent -= startOfPrevLine;

                if (indent > 0)
                {
                    queryField.value = e.newData.Insert(cursorIndex, new string(' ', indent));
                    int newCursorIndex = cursorIndex + indent;
                    queryField.SelectRange(newCursorIndex, newCursorIndex);
                }
            }
        }
    }

    private void queryField_KeyDown(KeyDownEvent e)
    {
        switch (e.keyCode)
        {
            case KeyCode.F5:
                runButton_Click();
                break;

            case KeyCode.Tab:
                string value = autoComplete.Count == 1 ? autoComplete[0].Substring(filterLength) : "    ";
                queryField.value = queryField.value.Insert(queryField.cursorIndex, value);
                int newCursorIndex = queryField.cursorIndex + value.Length;
                queryField.SelectRange(newCursorIndex, newCursorIndex);
                break;
        }
    }

    private void resultContainer_OnGUI()
    {
        UpdateAutoComplete();

        for (int i = 0; i < autoComplete.Count; i++)
            GUILayout.Label(autoComplete[i]);

        ShowResults(queryResult, position.height);
    }

    private void runButton_Click()
    {
        queryResult = CompileAndRunQuery(queryField.value, null);
    }

    private static Dictionary<string, Type> _allTypes;
    private static Dictionary<string, Type> allTypes
    {
        get
        {
            if (_allTypes == null)
            {
                _allTypes = loadedAssemblies
                    .SelectMany(i => i.GetTypes())
                    .Distinct(new SelectorEqualityComparer<Type, string>(i => i.Name))
                    .ToDictionary(i => i.Name, i => i);
            }
            return _allTypes;
        }
    }

    private static GUIStyle _boldLabel;
    private static GUIStyle boldLabel
    {
        get
        {
            if (_boldLabel == null)
            {
                _boldLabel = new GUIStyle(GUI.skin.label);
                _boldLabel.fontStyle = FontStyle.Bold;
            }
            return _boldLabel;
        }
    }

    private static GUIStyle _italicLabel;
    private static GUIStyle italicLabel
    {
        get
        {
            if (_italicLabel == null)
            {
                _italicLabel = new GUIStyle(GUI.skin.label);
                _italicLabel.fontStyle = FontStyle.Italic;
            }
            return _italicLabel;
        }
    }

    private static Assembly[] _loadedAssemblies;
    private static Assembly[] loadedAssemblies
    {
        get
        {
            if (_loadedAssemblies == null)
                _loadedAssemblies = AppDomain.CurrentDomain.GetAssemblies()
                    .Where(i => !i.IsDynamic)
                    .ToArray();
            return _loadedAssemblies;
        }
    }

    private VisualElement rootVisualElement
    {
        get
        {
            return GetType()
                .GetProperty("rootVisualContainer", BindingFlags.Instance | BindingFlags.NonPublic)
                .GetValue(this) as VisualElement;
        }
    }

    public static string _sourceTemplate;
    public static string sourceTemplate
    {
        get
        {
            if (_sourceTemplate == null)
            {
                var sb = new StringBuilder();

                foreach (var nameSpace in allTypes.Values
                    .Select(i => i.Namespace)
                    .Where(i => !string.IsNullOrEmpty(i) && !ignoredNamespaces.Any(j => i.StartsWith(j)))
                    .Distinct())
                {
                    sb.Append("using ").Append(nameSpace).Append(';').Append(Environment.NewLine);
                }

                sb.Append(@"
using Debug = UnityEngine.Debug;
using Object = UnityEngine.Object;

public static class Query
{{
    public static object Run()
    {{
        {0}
    }}
}}");

                _sourceTemplate = sb.ToString();
            }
            return _sourceTemplate;
        }
    }

    private static List<object> CompileAndRunQuery(string queryString, object[] parameters)
    {
        // It's important to use a fresh instance of compiler and parameters every time.
        var compiler = new CSharpCodeProvider();
        var compilerParams = new CompilerParameters();
        compilerParams.GenerateInMemory = true;
        foreach (var assembly in loadedAssemblies)
            compilerParams.ReferencedAssemblies.Add(assembly.Location);

        var compilationResult = compiler.CompileAssemblyFromSource(
            compilerParams, string.Format(sourceTemplate, queryString));

        var compilerErrors = compilationResult.Errors
            .Cast<CompilerError>()
            .Where(i => !i.IsWarning)
            .Cast<object>()
            .ToList();

        if (compilerErrors.Count > 0)
            return compilerErrors;

        var method = compilationResult.CompiledAssembly.GetType("Query").GetMethod("Run");
        try
        {
            var methodResult = method.Invoke(null, parameters);
            var enumerable = methodResult as IEnumerable;
            if (enumerable != null && !(methodResult is string))
            {
                return enumerable
                    .Cast<object>()
                    .ToList();
            }
            else
            {
                return new List<object>() { methodResult };
            }
        }
        catch (Exception ex)
        {
            var exceptions = new List<object>();
            while (ex != null)
            {
                exceptions.Add(ex);
                ex = ex.InnerException;
            }
            return exceptions;
        }
    }

    // Try to guess the type of an incomplete expression using regular expressions and reflection.
    private static TypeGuess GuessType(string code)
    {
        BindingFlags bindingFlags = BindingFlags.FlattenHierarchy | BindingFlags.Public;
        Match match;

        // Type
        match = Regex.Match(code, @"([A-Z]\w+)\.$");
        if (match.Success)
        {
            bindingFlags |= BindingFlags.Static;
            string typeName = match.Groups[1].Value;

            allTypes.TryGetValue(typeName, out Type type);
            return new TypeGuess(type, bindingFlags);
        }

        // Field or Property
        match = Regex.Match(code, @"^(.*?\.)(\w+)\.$");
        if (match.Success)
        {
            bindingFlags |= BindingFlags.Instance;
            string leftCode = match.Groups[1].Value;
            string memberName = match.Groups[2].Value;

            var leftType = GuessType(leftCode);
            if (leftType.Type == null)
                return new TypeGuess(null, bindingFlags);

            MemberInfo[] member = leftType.Type.GetMember(memberName,
                MemberTypes.Field | MemberTypes.Property, bindingFlags);
            if (member.Length == 0)
                return new TypeGuess(null, bindingFlags);

            var property = member[0] as PropertyInfo;
            if (property != null)
                return new TypeGuess(property.PropertyType, bindingFlags);

            var field = member[0] as FieldInfo;
            if (field != null)
                return new TypeGuess(field.FieldType, bindingFlags);

            return new TypeGuess(null, bindingFlags);
        }

        // Method
        match = Regex.Match(code, @"^(.*?)(\w+)\([^)]*\)+\.$");
        if (match.Success)
        {
            bindingFlags |= BindingFlags.Instance;
            string leftCode = match.Groups[1].Value;
            string methodName = match.Groups[2].Value;

            var leftType = GuessType(leftCode);
            if (leftType.Type == null)
                return new TypeGuess(null, bindingFlags);

            MethodInfo method = leftType.Type.GetMethod(methodName);
            if (method == null)
                return new TypeGuess(null, bindingFlags);

            Type returnType = method.ReturnType;
            return new TypeGuess(returnType, bindingFlags);
        }

        // Generic Method
        match = Regex.Match(code, @"\.\w+<(\w+)>\([^)]*\)+\.$");
        if (match.Success)
        {
            bindingFlags |= BindingFlags.Instance;
            string genericTypeName = match.Groups[1].Value;

            allTypes.TryGetValue(genericTypeName, out Type genericType);
            return new TypeGuess(genericType, bindingFlags);
        }

        // LINQ
        match = Regex.Match(code, @"^(.*?\.)\w+\((\w+)\s*=>.*?(\w+)\.$");
        if (match.Success && match.Groups[2].Value == match.Groups[3].Value)
        {
            string leftCode = match.Groups[1].Value;

            var typeLeft = GuessType(leftCode);
            return typeLeft;
        }

        return new TypeGuess(null, BindingFlags.Default);
    }

    private void UpdateAutoComplete()
    {
        autoComplete.Clear();
        filterLength = 0;

        if (queryField.cursorIndex <= 0)
            return;

        int leftDot = queryField.value.LastIndexOf('.', queryField.cursorIndex - 1);
        if (leftDot < 0)
            return;

        string code = queryField.value.Substring(0, leftDot + 1);
        filterLength = queryField.cursorIndex - leftDot - 1;

        var type = GuessType(code);
        if (type.Type == null)
            return;

        var members = type.Type.GetMembers(type.Flags);

        if (autoComplete.Capacity < members.Length)
            autoComplete.Capacity = members.Length;

        string filter = queryField.value.Substring(leftDot + 1, filterLength);

        for (int i = 0; i < members.Length; i++)
        {
            string name = members[i].Name;
            if ((filter == "" || name.StartsWith(filter)) && name[0] != '.' && !name.StartsWith("add_")
                && !name.StartsWith("get_") && !name.StartsWith("remove_") && !name.StartsWith("op_")
                && !name.StartsWith("set_"))
            {
                autoComplete.Add(members[i].Name);
            }
        }

        autoComplete.Sort();

        for (int i = autoComplete.Count - 1; i > 0; i--)
            if (autoComplete[i] == autoComplete[i - 1])
                autoComplete.RemoveAt(i);
    }

    private static void ShowResults(List<object> data, float maxHeight)
    {
        if (data == null)
        {
            GUILayout.Label("null", italicLabel);
            return;
        }

        if (data.Count == 0)
        {
            GUILayout.Label("No results", italicLabel);
            return;
        }

        int maxItems = Math.Min(data.Count, (int)(maxHeight / EditorGUIUtility.singleLineHeight));

        Type dataType = GetCommonBaseType(data
            .Where(i => i != null)
            .Select(i => i.GetType()));

        if (dataType == null)
        {
            for (int i = 0; i < maxItems; i++)
                GUILayout.Label("null", italicLabel);
            return;
        }

        if (typeof(CompilerError).IsAssignableFrom(dataType)
            || typeof(Exception).IsAssignableFrom(dataType))
        {
            for (int i = 0; i < maxItems; i++)
                GUILayout.Label(data[i].ToString());
            return;
        }

        var bindingFlags = BindingFlags.Instance | BindingFlags.Public | BindingFlags.FlattenHierarchy;

        var fields = dataType.GetFields(bindingFlags).ToList();

        var properties = dataType.GetProperties(bindingFlags)
            .Where(i => i.GetIndexParameters().Length == 0)
            .Where(i => i.GetCustomAttributes(typeof(ObsoleteAttribute), true).Length == 0)
            .Where(i => !ignoredProperties.Any(j => j.DeclaringType == i.DeclaringType && j.Name == i.Name))
            .ToList();

        int memberCount = fields.Count + properties.Count;

        if (memberCount == 0 || dataType == typeof(string))
        {
            for (int i = 0; i < maxItems; i++)
            {
                var item = data[i];
                if (item != null)
                    GUILayout.Label(item.ToString());
                else
                    GUILayout.Label("null", italicLabel);
            }

            return;
        }

        var columnWidth = GUILayout.Width(EditorGUIUtility.currentViewWidth / memberCount
                - EditorGUIUtility.standardVerticalSpacing);
        var rowHeight = GUILayout.Height(EditorGUIUtility.singleLineHeight);
        var options = new GUILayoutOption[] { columnWidth, rowHeight };

        GUILayout.BeginHorizontal();

        foreach (var field in fields)
            GUILayout.Toggle(false, new GUIContent(field.Name, field.FieldType.Name), boldLabel, options);

        foreach (var property in properties)
            GUILayout.Toggle(false, new GUIContent(property.Name, property.PropertyType.Name), boldLabel,
                options);

        GUILayout.EndHorizontal();

        for (int i = 0; i < maxItems; i++)
        {
            if (data[i] == null)
            {
                GUILayout.Label("null", italicLabel);
                continue;
            }

            GUILayout.BeginHorizontal();

            foreach (var field in fields)
            {
                try
                {
                    GridCell(field.GetValue(data[i]), options);
                }
                catch (Exception ex)
                {
                    GUILayout.Label(new GUIContent(ex.GetType().Name, ex.Message), italicLabel, options);
                }
            }

            foreach (var property in properties)
            {
                try
                {
                    GridCell(property.GetValue(data[i], null), options);
                }
                catch (Exception ex)
                {
                    GUILayout.Label(new GUIContent(ex.GetType().Name, ex.Message), italicLabel, options);
                }
            }

            GUILayout.EndHorizontal();
        }
    }

    private static Type GetCommonBaseType(IEnumerable<Type> types)
    {
        if (types == null || !types.Any())
            return null;

        var baseTypeLists = new List<List<Type>>();
        foreach (var type in types.Distinct())
        {
            var baseTypes = new List<Type>();
            Type baseType = type;
            while (baseType != null)
            {
                baseTypes.Add(baseType);
                baseType = baseType.BaseType;
            }
            baseTypes.Reverse();
            baseTypeLists.Add(baseTypes);
        }
        baseTypeLists.Sort((a, b) => a.Count - b.Count);

        Type commonAncestor = typeof(object);
        for (int i = 0; ; i++)
        {
            if (baseTypeLists[0].Count <= i)
                return commonAncestor;

            Type baseType = baseTypeLists[0][i];

            for (int j = 1; j < baseTypeLists.Count; j++)
                if (baseTypeLists[j][i] != baseType)
                    return commonAncestor;

            commonAncestor = baseType;
        }
    }

    private static void GridCell(object value, params GUILayoutOption[] options)
    {
        if (value == null)
            GUILayout.Label("null", italicLabel, options);
        else if (value is UnityEngine.Object)
            GridCell(value as UnityEngine.Object, options);
        else if (value is IEnumerable && !(value is string))
            GridCell(value as IEnumerable, options);
        else if (value is Matrix4x4)
            GridCell((Matrix4x4)value, options);
        else
            EditorGUILayout.SelectableLabel(Str(value), options);
    }

    private static void GridCell(Matrix4x4 value, params GUILayoutOption[] options)
    {
        var content = new GUIContent("", Str(value).Trim());

        if (value == Matrix4x4.identity) content.text = "identity";
        else if (value == Matrix4x4.zero) content.text = "zero";
        else if (value.m30 == 0 && value.m31 == 0 && value.m32 == 0) content.text = "affine";
        else content.text = "projection";

        GUILayout.Label(content, italicLabel, options);
    }

    private static void GridCell(IEnumerable value, params GUILayoutOption[] options)
    {
        var items = value
            .Cast<object>()
            .Take(10)
            .Select(i => Str(i))
            .ToArray();

        var content = new GUIContent();
        content.text = items.Length.ToString() + " item" + (items.Length != 1 ? "s" : "");
        content.tooltip = string.Join(Environment.NewLine, items);

        GUILayout.Label(content, italicLabel, options);
    }

    private static void GridCell(UnityEngine.Object value, params GUILayoutOption[] options)
    {
        Type valueType = value.GetType();
        var content = EditorGUIUtility.ObjectContent(value, valueType);
        content.text = Regex.Replace(content.text, @" \([^)]*\)$", "");
        content.tooltip = valueType.Name;

        if (GUILayout.Button(content, GUI.skin.label, options))
        {
            EditorGUIUtility.PingObject(value);
            if (Event.current.control || Event.current.shift)
            {
                var selection = Selection.objects;
                if (selection.Contains(value))
                    Selection.objects = selection
                        .Where(i => i != value)
                        .ToArray();
                else
                    Selection.objects = selection
                        .Concat(new[] { value })
                        .ToArray();
            }
            else
            {
                Selection.objects = new[] { value };
            }
        }
    }

    private static string Str(object value)
    {
        if (value == null) return null;
        else if (value is Quaternion) return Str((Quaternion)value);
        else if (value is Scene) return Str((Scene)value);
        else if (value is Vector2) return Str((Vector2)value);
        else if (value is Vector3) return Str((Vector3)value);
        else if (value is Vector4) return Str((Vector4)value);
        else return value.ToString();
    }

    private static string Str(Quaternion value)
    {
        return string.Format("{0}, {1}, {2}, {3}", value.x, value.y, value.z, value.w);
    }

    private static string Str(Scene value)
    {
        return value.name;
    }

    private static string Str(Vector2 value)
    {
        return string.Format("{0}, {1}", value.x, value.y);
    }

    private static string Str(Vector3 value)
    {
        return string.Format("{0}, {1}, {2}", value.x, value.y, value.z);
    }

    private static string Str(Vector4 value)
    {
        return string.Format("{0}, {1}, {2}, {3}", value.x, value.y, value.z, value.w);
    }

    private struct TypeGuess
    {
        public Type Type;
        public BindingFlags Flags;

        public TypeGuess(Type type, BindingFlags flags)
        {
            Type = type;
            Flags = flags;
        }
    }
}

public sealed class SelectorEqualityComparer<TSource, TResult> : IEqualityComparer<TSource>
{
    private Func<TSource, TResult> selector;

    public SelectorEqualityComparer(Func<TSource, TResult> selector)
    {
        this.selector = selector;
    }

    public bool Equals(TSource x, TSource y)
    {
        return EqualityComparer<TResult>.Default.Equals(selector(x), selector(y));
    }

    public int GetHashCode(TSource obj)
    {
        return selector(obj).GetHashCode();
    }
}