dev_treeparse.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### parsing functions\n",
    "using - https://github.com/tree-sitter-grammars/tree-sitter-glsl\n",
    "basic usage: https://github.com/tree-sitter/py-tree-sitter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tree_sitter_glsl as tsglsl\n",
    "import tree_sitter\n",
    "from tree_sitter import Language, Parser\n",
    "from typing import List, Tuple\n",
    "\n",
    "GLSL_LANGUAGE = Language(tsglsl.language())\n",
    "\n",
    "parser = Parser(GLSL_LANGUAGE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://gist.github.com/TACIXAT/c5b2db4a80c812c4b4373b65e179a220\n",
    "def format_sexpression(s, indent_level=0, indent_size=4):\n",
    "    \"\"\"ChatGPT + TACIXAT\"\"\"\n",
    "    output = \"\"\n",
    "    i = 0\n",
    "    # Initialize to False to avoid newline for the first token\n",
    "    need_newline = False\n",
    "    cdepth = [] # Track colons\n",
    "    while i < len(s):\n",
    "        if s[i] == \"(\":\n",
    "            output += \"\\n\" + \" \" * (indent_level * indent_size) + \"(\"\n",
    "            indent_level += 1\n",
    "            need_newline = False  # Avoid newline after opening parenthesis\n",
    "        elif s[i] == \":\":\n",
    "            indent_level += 1\n",
    "            cdepth.append(indent_level) # Store depth where we saw colon\n",
    "            output += \":\"\n",
    "        elif s[i] == \")\":\n",
    "            indent_level -= 1\n",
    "            if len(cdepth) > 0 and indent_level == cdepth[-1]: \n",
    "                # Unindent when we return to the depth we saw the last colon\n",
    "                cdepth.pop()\n",
    "                indent_level -= 1\n",
    "            output += \")\"\n",
    "            need_newline = True  # Newline needed after closing parenthesis\n",
    "        elif s[i] == \" \":\n",
    "            output += \" \"\n",
    "        else:\n",
    "            j = i\n",
    "            while j < len(s) and s[j] not in [\"(\", \")\", \" \", \":\"]:\n",
    "                j += 1\n",
    "            # Add newline and indentation only when needed\n",
    "            if need_newline:\n",
    "                output += \"\\n\" + \" \" * (indent_level * indent_size)\n",
    "            output += s[i:j]\n",
    "            i = j - 1\n",
    "            need_newline = True  # Next token should start on a new line\n",
    "        i += 1\n",
    "    return output\n",
    "\n",
    "def color_ranges(code:str, func_bytes:list):\n",
    "    \"\"\"\n",
    "    returns a string with ANSI color codes to color parts of functions. Omits some part of the code!\n",
    "    \"\"\"\n",
    "    #TODO: don't skip original code (return as default colored) \n",
    "    code_bytes = bytes(code, encoding=\"utf-8\")\n",
    "    colored_bytes = b\"\\x1b[0m\"\n",
    "    for start_comment, start_header, end_header, end_docstring, end_function in func_bytes:\n",
    "        colored_bytes += b\"\\x1b[32m\" + code_bytes[start_comment:start_header]\n",
    "        colored_bytes += b\"\\x1b[31m\" + code_bytes[start_header:end_header]\n",
    "        colored_bytes += b\"\\x1b[33m\" + code_bytes[end_header:end_docstring]\n",
    "        colored_bytes += b\"\\x1b[34m\" + code_bytes[end_docstring:end_function]\n",
    "        colored_bytes += b\"\\x1b[0m\" #escape back to normal\n",
    "        colored_bytes += b\"\\n\\n\" # for pretty but also lazy.\n",
    "\n",
    "    return colored_bytes.decode(encoding=\"utf-8\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tree_sitter.Tree at 0x191eddc3f50>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "example_code = \"\"\"//glsl\n",
    "// comment at the start specifying that we use Apache 2.0 license :)\n",
    "\n",
    "\n",
    "\n",
    "/*\n",
    "*what about multi line comments?\n",
    "*\n",
    "*that use this special delimiter?\n",
    "*\n",
    "*/\n",
    "vec3 red(float intensity){\n",
    "    // comment inside the function body, but the top\n",
    "    intensity = max(intensity, 1.0);\n",
    "\n",
    "    // comment deeper inside the function body\n",
    "    return vec3(1.0,0.0,0.0) * intensity;\n",
    "}\n",
    "\n",
    "// a comment between functions where there is nothing more.\n",
    "\n",
    "vec3 green(float intensity){\n",
    "    // docstring inside\n",
    "    // second inside docstring\n",
    "    intensity = max(intensity, 1.0);\n",
    "\n",
    "    return vec3(0.0,1.0,0.0) * intensity;\n",
    "}\n",
    "\n",
    "// this comment is somewhere in between we don't want it\n",
    "\n",
    "\n",
    "// comment one before\n",
    "// comment two before\n",
    "vec3 blue(float intensity){\n",
    "    intensity = max(intensity, 1.0);\n",
    "\n",
    "    return vec3(0.0,0.0,1.0) * intensity;\n",
    "}\n",
    "\n",
    "void mainImage( out vec4 fragColor, in vec2 fragCoord )\n",
    "{\n",
    "    vec2 uv = fragCoord/iResolution.xy;\n",
    "\n",
    "    vec3 col = mix(red(0.5), blue(0.3), 0.8);\n",
    "    col = mix(col, green(1.0), sin(iTime));\n",
    "    fragColor = vec4(col,1.0);\n",
    "}\n",
    "\n",
    "// rear comment\n",
    "\"\"\"\n",
    "\n",
    "tree = parser.parse(bytes(example_code, encoding=\"utf-8\"))\n",
    "root_node = tree.root_node\n",
    "tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "165\n",
      "153\n",
      "153\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[b'vec3',\n",
       " b'id',\n",
       " b'(',\n",
       " b'float',\n",
       " b'id',\n",
       " b')',\n",
       " b'{',\n",
       " b'id',\n",
       " b'=',\n",
       " b'id',\n",
       " b'(',\n",
       " b'id',\n",
       " b',',\n",
       " b'1.0',\n",
       " b')',\n",
       " b';',\n",
       " b'return',\n",
       " b'id',\n",
       " b'(',\n",
       " b'1.0',\n",
       " b',',\n",
       " b'0.0',\n",
       " b',',\n",
       " b'0.0',\n",
       " b')',\n",
       " b'*',\n",
       " b'id',\n",
       " b';',\n",
       " b'}',\n",
       " b'vec3',\n",
       " b'id',\n",
       " b'(',\n",
       " b'float',\n",
       " b'id',\n",
       " b')',\n",
       " b'{',\n",
       " b'id',\n",
       " b'=',\n",
       " b'id',\n",
       " b'(',\n",
       " b'id',\n",
       " b',',\n",
       " b'1.0',\n",
       " b')',\n",
       " b';',\n",
       " b'return',\n",
       " b'id',\n",
       " b'(',\n",
       " b'0.0',\n",
       " b',',\n",
       " b'1.0',\n",
       " b',',\n",
       " b'0.0',\n",
       " b')',\n",
       " b'*',\n",
       " b'id',\n",
       " b';',\n",
       " b'}',\n",
       " b'vec3',\n",
       " b'id',\n",
       " b'(',\n",
       " b'float',\n",
       " b'id',\n",
       " b')',\n",
       " b'{',\n",
       " b'id',\n",
       " b'=',\n",
       " b'id',\n",
       " b'(',\n",
       " b'id',\n",
       " b',',\n",
       " b'1.0',\n",
       " b')',\n",
       " b';',\n",
       " b'return',\n",
       " b'id',\n",
       " b'(',\n",
       " b'0.0',\n",
       " b',',\n",
       " b'0.0',\n",
       " b',',\n",
       " b'1.0',\n",
       " b')',\n",
       " b'*',\n",
       " b'id',\n",
       " b';',\n",
       " b'}',\n",
       " b'void',\n",
       " b'id',\n",
       " b'(',\n",
       " b'out',\n",
       " b'vec4',\n",
       " b'id',\n",
       " b',',\n",
       " b'in',\n",
       " b'vec2',\n",
       " b'id',\n",
       " b')',\n",
       " b'{',\n",
       " b'vec2',\n",
       " b'id',\n",
       " b'=',\n",
       " b'id',\n",
       " b'/',\n",
       " b'id',\n",
       " b'.',\n",
       " b'xy',\n",
       " b';',\n",
       " b'vec3',\n",
       " b'id',\n",
       " b'=',\n",
       " b'id',\n",
       " b'(',\n",
       " b'id',\n",
       " b'(',\n",
       " b'0.5',\n",
       " b')',\n",
       " b',',\n",
       " b'id',\n",
       " b'(',\n",
       " b'0.3',\n",
       " b')',\n",
       " b',',\n",
       " b'0.8',\n",
       " b')',\n",
       " b';',\n",
       " b'id',\n",
       " b'=',\n",
       " b'id',\n",
       " b'(',\n",
       " b'id',\n",
       " b',',\n",
       " b'id',\n",
       " b'(',\n",
       " b'1.0',\n",
       " b')',\n",
       " b',',\n",
       " b'id',\n",
       " b'(',\n",
       " b'id',\n",
       " b')',\n",
       " b')',\n",
       " b';',\n",
       " b'id',\n",
       " b'=',\n",
       " b'id',\n",
       " b'(',\n",
       " b'id',\n",
       " b',',\n",
       " b'1.0',\n",
       " b')',\n",
       " b';',\n",
       " b'}']"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# simple lexer?\n",
    "\n",
    "def get_leaves(subtree: tree_sitter.Node, skip_comments=False, rename_identifiers=False):\n",
    "    # TODO: get like a wrapper function to give the root node initially...\n",
    "    tokens = []\n",
    "    if subtree.child_count == 0:\n",
    "        if subtree.type == \"comment\" and skip_comments:\n",
    "            pass\n",
    "        elif subtree.type == \"identifier\" and rename_identifiers:\n",
    "            # TODO: what about different name spaces - where do we hand this upwards?\n",
    "            # do we need to like return our mapping?\n",
    "            return [f\"id\"]\n",
    "        else:\n",
    "            return [subtree.text]\n",
    "    else:\n",
    "        for child in subtree.children:\n",
    "            tokens.extend(get_leaves(child, skip_comments, rename_identifiers))\n",
    "    return tokens\n",
    "\n",
    "exact_tokens = get_leaves(tree.root_node, skip_comments=False)\n",
    "lexical_tokens = get_leaves(tree.root_node, skip_comments=True)\n",
    "syntax_tokens = get_leaves(tree.root_node, True, True)\n",
    "print(len(exact_tokens))\n",
    "print(len(lexical_tokens))\n",
    "print(len(syntax_tokens))\n",
    "\n",
    "# tree.root_node.children[3].children[1].children[0].text\n",
    "syntax_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "(translation_unit \n",
      "(comment) \n",
      "(comment) \n",
      "(comment) \n",
      "(function_definition \n",
      "    type: \n",
      "        (type_identifier) \n",
      "    declarator: \n",
      "        (function_declarator \n",
      "            declarator: \n",
      "                (identifier) \n",
      "            parameters: \n",
      "                (parameter_list \n",
      "                    (parameter_declaration \n",
      "                        type: \n",
      "                            (primitive_type) \n",
      "                        declarator: \n",
      "                            (identifier)))) \n",
      "    body: \n",
      "        (compound_statement \n",
      "            (comment) \n",
      "            (expression_statement \n",
      "                (assignment_expression \n",
      "                    left: \n",
      "                        (identifier) \n",
      "                    right: \n",
      "                        (call_expression \n",
      "                            function: \n",
      "                                (identifier) \n",
      "                            arguments: \n",
      "                                (argument_list \n",
      "                                    (identifier) \n",
      "                                    (number_literal))))) \n",
      "            (comment) \n",
      "            (return_statement \n",
      "                (binary_expression \n",
      "                    left: \n",
      "                        (call_expression \n",
      "                            function: \n",
      "                                (identifier) \n",
      "                            arguments: \n",
      "                                (argument_list \n",
      "                                    (number_literal) \n",
      "                                    (number_literal) \n",
      "                                    (number_literal))) \n",
      "                    right: \n",
      "                        (identifier))))) \n",
      "(comment) \n",
      "(function_definition \n",
      "    type: \n",
      "        (type_identifier) \n",
      "    declarator: \n",
      "        (function_declarator \n",
      "            declarator: \n",
      "                (identifier) \n",
      "            parameters: \n",
      "                (parameter_list \n",
      "                    (parameter_declaration \n",
      "                        type: \n",
      "                            (primitive_type) \n",
      "                        declarator: \n",
      "                            (identifier)))) \n",
      "    body: \n",
      "        (compound_statement \n",
      "            (comment) \n",
      "            (comment) \n",
      "            (expression_statement \n",
      "                (assignment_expression \n",
      "                    left: \n",
      "                        (identifier) \n",
      "                    right: \n",
      "                        (call_expression \n",
      "                            function: \n",
      "                                (identifier) \n",
      "                            arguments: \n",
      "                                (argument_list \n",
      "                                    (identifier) \n",
      "                                    (number_literal))))) \n",
      "            (return_statement \n",
      "                (binary_expression \n",
      "                    left: \n",
      "                        (call_expression \n",
      "                            function: \n",
      "                                (identifier) \n",
      "                            arguments: \n",
      "                                (argument_list \n",
      "                                    (number_literal) \n",
      "                                    (number_literal) \n",
      "                                    (number_literal))) \n",
      "                    right: \n",
      "                        (identifier))))) \n",
      "(comment) \n",
      "(comment) \n",
      "(comment) \n",
      "(function_definition \n",
      "    type: \n",
      "        (type_identifier) \n",
      "    declarator: \n",
      "        (function_declarator \n",
      "            declarator: \n",
      "                (identifier) \n",
      "            parameters: \n",
      "                (parameter_list \n",
      "                    (parameter_declaration \n",
      "                        type: \n",
      "                            (primitive_type) \n",
      "                        declarator: \n",
      "                            (identifier)))) \n",
      "    body: \n",
      "        (compound_statement \n",
      "            (expression_statement \n",
      "                (assignment_expression \n",
      "                    left: \n",
      "                        (identifier) \n",
      "                    right: \n",
      "                        (call_expression \n",
      "                            function: \n",
      "                                (identifier) \n",
      "                            arguments: \n",
      "                                (argument_list \n",
      "                                    (identifier) \n",
      "                                    (number_literal))))) \n",
      "            (return_statement \n",
      "                (binary_expression \n",
      "                    left: \n",
      "                        (call_expression \n",
      "                            function: \n",
      "                                (identifier) \n",
      "                            arguments: \n",
      "                                (argument_list \n",
      "                                    (number_literal) \n",
      "                                    (number_literal) \n",
      "                                    (number_literal))) \n",
      "                    right: \n",
      "                        (identifier))))) \n",
      "(function_definition \n",
      "    type: \n",
      "        (primitive_type) \n",
      "    declarator: \n",
      "        (function_declarator \n",
      "            declarator: \n",
      "                (identifier) \n",
      "            parameters: \n",
      "                (parameter_list \n",
      "                    (parameter_declaration \n",
      "                        type: \n",
      "                            (type_identifier) \n",
      "                        declarator: \n",
      "                            (identifier)) \n",
      "                    (parameter_declaration \n",
      "                        type: \n",
      "                            (type_identifier) \n",
      "                        declarator: \n",
      "                            (identifier)))) \n",
      "    body: \n",
      "        (compound_statement \n",
      "            (declaration \n",
      "                type: \n",
      "                    (type_identifier) \n",
      "                declarator: \n",
      "                    (init_declarator \n",
      "                        declarator: \n",
      "                            (identifier) \n",
      "                        value: \n",
      "                            (binary_expression \n",
      "                                left: \n",
      "                                    (identifier) \n",
      "                                right: \n",
      "                                    (field_expression \n",
      "                                        argument: \n",
      "                                            (identifier) \n",
      "                                        field: \n",
      "                                            (field_identifier))))) \n",
      "            (declaration \n",
      "                type: \n",
      "                    (type_identifier) \n",
      "                declarator: \n",
      "                    (init_declarator \n",
      "                        declarator: \n",
      "                            (identifier) \n",
      "                        value: \n",
      "                            (call_expression \n",
      "                                function: \n",
      "                                    (identifier) \n",
      "                                arguments: \n",
      "                                    (argument_list \n",
      "                                        (call_expression \n",
      "                                            function: \n",
      "                                                (identifier) \n",
      "                                            arguments: \n",
      "                                                (argument_list \n",
      "                                                    (number_literal))) \n",
      "                                        (call_expression \n",
      "                                            function: \n",
      "                                                (identifier) \n",
      "                                            arguments: \n",
      "                                                (argument_list \n",
      "                                                    (number_literal))) \n",
      "                                        (number_literal))))) \n",
      "            (expression_statement \n",
      "                (assignment_expression \n",
      "                    left: \n",
      "                        (identifier) \n",
      "                    right: \n",
      "                        (call_expression \n",
      "                            function: \n",
      "                                (identifier) \n",
      "                            arguments: \n",
      "                                (argument_list \n",
      "                                    (identifier) \n",
      "                                    (call_expression \n",
      "                                        function: \n",
      "                                            (identifier) \n",
      "                                        arguments: \n",
      "                                            (argument_list \n",
      "                                                (number_literal))) \n",
      "                                    (call_expression \n",
      "                                        function: \n",
      "                                            (identifier) \n",
      "                                        arguments: \n",
      "                                            (argument_list \n",
      "                                                (identifier))))))) \n",
      "            (expression_statement \n",
      "                (assignment_expression \n",
      "                    left: \n",
      "                        (identifier) \n",
      "                    right: \n",
      "                        (call_expression \n",
      "                            function: \n",
      "                                (identifier) \n",
      "                            arguments: \n",
      "                                (argument_list \n",
      "                                    (identifier) \n",
      "                                    (number_literal))))))))\n"
     ]
    }
   ],
   "source": [
    "print(format_sexpression(str(root_node), -1, 4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2\n",
      "(0, {'comment.before': [<Node type=comment, start_point=(1, 0), end_point=(1, 68)>, <Node type=comment, start_point=(5, 0), end_point=(5, 41)>], 'function.type': <Node type=primitive_type, start_point=(33, 0), end_point=(33, 4)>, 'function.declarator': <Node type=function_declarator, start_point=(33, 5), end_point=(33, 55)>, 'function.body': <Node type=compound_statement, start_point=(34, 0), end_point=(40, 1)>})\n",
      "{'comment.before': [<Node type=comment, start_point=(1, 0), end_point=(1, 68)>, <Node type=comment, start_point=(5, 0), end_point=(5, 41)>], 'function.type': <Node type=primitive_type, start_point=(33, 0), end_point=(33, 4)>, 'function.declarator': <Node type=function_declarator, start_point=(33, 5), end_point=(33, 55)>, 'function.body': <Node type=compound_statement, start_point=(34, 0), end_point=(40, 1)>}\n"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "'list' object has no attribute 'start_byte'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[5], line 53\u001b[0m\n\u001b[0;32m     50\u001b[0m         funcs\u001b[38;5;241m.\u001b[39mappend(\u001b[38;5;28mtuple\u001b[39m([start_comment, start_header, end_header, end_docstring, end_function]))\n\u001b[0;32m     51\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m funcs\n\u001b[1;32m---> 53\u001b[0m funcs2 \u001b[38;5;241m=\u001b[39m \u001b[43mparse_functions2\u001b[49m\u001b[43m(\u001b[49m\u001b[43mexample_code\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     54\u001b[0m \u001b[38;5;28mprint\u001b[39m(funcs2)\n\u001b[0;32m     55\u001b[0m colored_funcs2 \u001b[38;5;241m=\u001b[39m color_ranges(example_code, func_bytes\u001b[38;5;241m=\u001b[39mfuncs2)\n",
      "Cell \u001b[1;32mIn[5], line 46\u001b[0m, in \u001b[0;36mparse_functions2\u001b[1;34m(code)\u001b[0m\n\u001b[0;32m     43\u001b[0m end_docstring \u001b[38;5;241m=\u001b[39m end_header \u001b[38;5;241m=\u001b[39m m[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfunction.declarator\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mend_byte\n\u001b[0;32m     44\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcomment.before\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m m:\n\u001b[0;32m     45\u001b[0m     \u001b[38;5;66;03m# print(len(m[\"comment.before\"]))\u001b[39;00m\n\u001b[1;32m---> 46\u001b[0m     start_comment \u001b[38;5;241m=\u001b[39m \u001b[43mm\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mcomment.before\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstart_byte\u001b[49m\n\u001b[0;32m     47\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdocstring\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m m:\n\u001b[0;32m     48\u001b[0m     end_docstring \u001b[38;5;241m=\u001b[39m m[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdocstring\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mend_byte\n",
      "\u001b[1;31mAttributeError\u001b[0m: 'list' object has no attribute 'start_byte'"
     ]
    }
   ],
   "source": [
    "# https://tree-sitter.github.io/tree-sitter/using-parsers#query-syntax\n",
    "\n",
    "func_query = GLSL_LANGUAGE.query(\n",
    "        \"\"\"\n",
    "        (function_definition ) @function\n",
    "    \"\"\"\n",
    "    )\n",
    "\n",
    "comment_query = GLSL_LANGUAGE.query(\n",
    "    \"\"\"\n",
    "    (\n",
    "        (comment)* @comment.before\n",
    "        (function_definition \n",
    "            .\n",
    "            type: (primitive_type) @function.type\n",
    "            declarator: (function_declarator) @function.declarator\n",
    "            body: (compound_statement\n",
    "                (comment)? @docstring\n",
    "                ) @function.body\n",
    "        )\n",
    "    )\n",
    "    \"\"\"\n",
    ")\n",
    "\n",
    "def parse_functions2(code:str) -> List[Tuple[int,int,int,int,int]]:\n",
    "    tree = parser.parse(bytes(code, encoding=\"utf-8\"))\n",
    "    root_node = tree.root_node\n",
    "    funcs = []\n",
    "    \n",
    "    matches = comment_query.matches(root_node, end_byte=500)\n",
    "    print(len(matches))\n",
    "    for func in matches:\n",
    "        # func = func[1][\"function\"]\n",
    "        # print(func.end_point)\n",
    "        # matches = comment_query.matches(root_node, end_point=func.end_point)\n",
    "        print(func)\n",
    "        # continue\n",
    "        m = matches[0][1]\n",
    "        print(m)\n",
    "        start_comment = start_header = end_header = end_docstring = end_function = None\n",
    "        start_comment = start_header = m[\"function.type\"].start_byte\n",
    "        end_function = m[\"function.body\"].end_byte\n",
    "        end_docstring = end_header = m[\"function.declarator\"].end_byte\n",
    "        if \"comment.before\" in m:\n",
    "            # print(len(m[\"comment.before\"]))\n",
    "            start_comment = m[\"comment.before\"].start_byte\n",
    "        if \"docstring\" in m:\n",
    "            end_docstring = m[\"docstring\"].end_byte\n",
    "        \n",
    "        funcs.append(tuple([start_comment, start_header, end_header, end_docstring, end_function]))\n",
    "    return funcs\n",
    "\n",
    "funcs2 = parse_functions2(example_code)\n",
    "print(funcs2)\n",
    "colored_funcs2 = color_ranges(example_code, func_bytes=funcs2)\n",
    "# print(colored_funcs2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "20\n",
      "(<Node type=comment, start_point=(1, 0), end_point=(1, 68)>, 'comment')\n",
      "(<Node type=comment, start_point=(1, 0), end_point=(1, 68)>, 'comment')\n",
      "(<Node type=comment, start_point=(1, 0), end_point=(1, 68)>, 'comment')\n",
      "(<Node type=comment, start_point=(1, 0), end_point=(1, 68)>, 'comment')\n",
      "(<Node type=comment, start_point=(5, 0), end_point=(5, 41)>, 'comment')\n",
      "(<Node type=comment, start_point=(5, 0), end_point=(5, 41)>, 'comment')\n",
      "(<Node type=comment, start_point=(5, 0), end_point=(5, 41)>, 'comment')\n",
      "(<Node type=comment, start_point=(5, 0), end_point=(5, 41)>, 'comment')\n",
      "(<Node type=function_definition, start_point=(6, 0), end_point=(12, 1)>, 'function')\n",
      "(<Node type=function_definition, start_point=(14, 0), end_point=(20, 1)>, 'function')\n",
      "(<Node type=comment, start_point=(22, 0), end_point=(22, 56)>, 'comment')\n",
      "(<Node type=comment, start_point=(22, 0), end_point=(22, 56)>, 'comment')\n",
      "(<Node type=comment, start_point=(25, 0), end_point=(25, 21)>, 'comment')\n",
      "(<Node type=comment, start_point=(25, 0), end_point=(25, 21)>, 'comment')\n",
      "(<Node type=comment, start_point=(26, 0), end_point=(26, 21)>, 'comment')\n",
      "(<Node type=comment, start_point=(26, 0), end_point=(26, 21)>, 'comment')\n",
      "(<Node type=function_definition, start_point=(27, 0), end_point=(31, 1)>, 'function')\n",
      "(<Node type=function_definition, start_point=(27, 0), end_point=(31, 1)>, 'function')\n",
      "(<Node type=function_definition, start_point=(33, 0), end_point=(40, 1)>, 'function')\n",
      "(<Node type=function_definition, start_point=(33, 0), end_point=(40, 1)>, 'function')\n"
     ]
    }
   ],
   "source": [
    "dev_query = GLSL_LANGUAGE.query(\n",
    "    \"\"\"\n",
    "    (   \n",
    "        [\n",
    "        (comment)\n",
    "        (comment)+\n",
    "        ] @comment\n",
    "        (function_definition ) @function\n",
    "    )\n",
    "    \"\"\"\n",
    ")\n",
    "\n",
    "ms = dev_query.captures(root_node)\n",
    "print(len(ms))\n",
    "for c in ms:\n",
    "    print(c)\n",
    "    # print(c[1][\"commented_function\"].text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(79, 156, 182, 235, 364), (427, 427, 455, 510, 592), (653, 697, 724, 724, 806), (808, 808, 865, 865, 1029)]\n",
      "\u001b[0m\u001b[32m/*\n",
      "*what about multi line comments?\n",
      "*\n",
      "*that use this special delimiter?\n",
      "*\n",
      "*/\n",
      "\u001b[31mvec3 red(float intensity){\u001b[33m\n",
      "    // comment inside the function body, but the top\u001b[34m\n",
      "    intensity = max(intensity, 1.0);\n",
      "\n",
      "    // comment deeper inside the function body\n",
      "    return vec3(1.0,0.0,0.0) * intensity;\n",
      "}\u001b[0m\n",
      "\n",
      "\u001b[32m\u001b[31mvec3 green(float intensity){\u001b[33m\n",
      "    // docstring inside\n",
      "    // second inside docstring\u001b[34m\n",
      "    intensity = max(intensity, 1.0);\n",
      "\n",
      "    return vec3(0.0,1.0,0.0) * intensity;\n",
      "}\u001b[0m\n",
      "\n",
      "\u001b[32m// comment one before\n",
      "// comment two before\n",
      "\u001b[31mvec3 blue(float intensity){\u001b[33m\u001b[34m\n",
      "    intensity = max(intensity, 1.0);\n",
      "\n",
      "    return vec3(0.0,0.0,1.0) * intensity;\n",
      "}\u001b[0m\n",
      "\n",
      "\u001b[32m\u001b[31mvoid mainImage( out vec4 fragColor, in vec2 fragCoord )\n",
      "{\u001b[33m\u001b[34m\n",
      "    vec2 uv = fragCoord/iResolution.xy;\n",
      "\n",
      "    vec3 col = mix(red(0.5), blue(0.3), 0.8);\n",
      "    col = mix(col, green(1.0), sin(iTime));\n",
      "    fragColor = vec4(col,1.0);\n",
      "}\u001b[0m\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from annotate import parse_functions as parse_functions_imported\n",
    "funcs = parse_functions_imported(example_code)\n",
    "print(funcs)\n",
    "colored_funcs = color_ranges(example_code, func_bytes=funcs)\n",
    "print(colored_funcs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(79, 156, 182, 235, 364), (366, 427, 455, 510, 592), (653, 697, 724, 724, 806), (808, 808, 865, 865, 1029)]\n",
      "\u001b[0m\u001b[32m/*\n",
      "*what about multi line comments?\n",
      "*\n",
      "*that use this special delimiter?\n",
      "*\n",
      "*/\n",
      "\u001b[31mvec3 red(float intensity){\u001b[33m\n",
      "    // comment inside the function body, but the top\u001b[34m\n",
      "    intensity = max(intensity, 1.0);\n",
      "\n",
      "    // comment deeper inside the function body\n",
      "    return vec3(1.0,0.0,0.0) * intensity;\n",
      "}\u001b[0m\n",
      "\n",
      "\u001b[32m// a comment between functions where there is nothing more.\n",
      "\n",
      "\u001b[31mvec3 green(float intensity){\u001b[33m\n",
      "    // docstring inside\n",
      "    // second inside docstring\u001b[34m\n",
      "    intensity = max(intensity, 1.0);\n",
      "\n",
      "    return vec3(0.0,1.0,0.0) * intensity;\n",
      "}\u001b[0m\n",
      "\n",
      "\u001b[32m// comment one before\n",
      "// comment two before\n",
      "\u001b[31mvec3 blue(float intensity){\u001b[33m\u001b[34m\n",
      "    intensity = max(intensity, 1.0);\n",
      "\n",
      "    return vec3(0.0,0.0,1.0) * intensity;\n",
      "}\u001b[0m\n",
      "\n",
      "\u001b[32m\u001b[31mvoid mainImage( out vec4 fragColor, in vec2 fragCoord )\n",
      "{\u001b[33m\u001b[34m\n",
      "    vec2 uv = fragCoord/iResolution.xy;\n",
      "\n",
      "    vec3 col = mix(red(0.5), blue(0.3), 0.8);\n",
      "    col = mix(col, green(1.0), sin(iTime));\n",
      "    fragColor = vec4(col,1.0);\n",
      "}\u001b[0m\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "\n",
    "def parse_functions(code:str) -> List[Tuple[int,int,int,int,int]]:\n",
    "    \"\"\"\n",
    "    parses the code using tree-parser-glsl\n",
    "    returns the **byte-indecies** for before_comment, start header, end header, end docstring, end_function.\n",
    "    returns a list 5-tupel. If before_comment or docstring aren't found, the indiecies will coinside with the next one.\n",
    "    \"\"\"\n",
    "    tree = parser.parse(bytes(code, encoding=\"utf-8\"))\n",
    "    root_node = tree.root_node\n",
    "    funcs = []\n",
    "    \n",
    "    # lazy init\n",
    "    start_comment = start_header = end_header = end_docstring = end_function = None\n",
    "    comment_line = -2\n",
    "    for child in root_node.children:\n",
    "        if (child.type == \"comment\"):\n",
    "            if ((comment_line + 1) != child.start_point[0]):\n",
    "                start_comment = child.start_byte\n",
    "            comment_line = child.end_point[0]\n",
    "        elif child.type == \"function_definition\":\n",
    "            start_header = child.start_byte\n",
    "            if not start_comment:\n",
    "                start_comment = start_header\n",
    "            end_function = child.end_byte\n",
    "            end_header = child.children[-1].children[0].end_byte\n",
    "            # inside the function body, past the \"{\"\n",
    "            for sub_child in child.children[-1].children[1:]:\n",
    "                if sub_child.type == \"comment\":\n",
    "                    end_docstring = sub_child.end_byte\n",
    "                else:\n",
    "                    if not end_docstring:\n",
    "                        end_docstring = end_header\n",
    "                    break\n",
    "                \n",
    "\n",
    "\n",
    "            funcs.append(tuple([start_comment, start_header, end_header, end_docstring, end_function]))\n",
    "            start_comment = start_header = end_header = end_docstring = end_function = None\n",
    "    return funcs\n",
    "\n",
    "\n",
    "funcs = parse_functions(example_code)\n",
    "print(funcs)\n",
    "colored_funcs = color_ranges(example_code, func_bytes=funcs)\n",
    "print(colored_funcs)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}