From dad30fc7cfbb50fb86e2652ef6a1afa2b38d8f8c Mon Sep 17 00:00:00 2001 From: marco-dossantos Date: Thu, 31 Jul 2025 12:13:59 +0200 Subject: [PATCH] feat(infotree): add extract_data processing and tests --- client/kimina_client/infotree.py | 1604 +++++++++++++++-------------- client/kimina_client/utils.py | 44 + setup.sh | 6 + tests/test_infotree_processing.py | 18 + 4 files changed, 884 insertions(+), 788 deletions(-) create mode 100644 tests/test_infotree_processing.py diff --git a/client/kimina_client/infotree.py b/client/kimina_client/infotree.py index 4f9156b..b73c5ed 100644 --- a/client/kimina_client/infotree.py +++ b/client/kimina_client/infotree.py @@ -1,788 +1,816 @@ -# import re -# from typing import Optional - - -# def extract_nodes_and_edges( -# infotree: list[dict], -# parent_id: Optional[int] = None, -# start_id: int = 0, -# include_failed_pp: bool = True, -# deduplicate: bool = False, -# ): -# """ -# Recursively extract nodes and edges from an infotree. - -# Parameters -# ---------- -# infotree : list of dict -# A list of dictionaries each containing 'node' and optionally 'children'. -# parent_id : int or None, optional -# The ID of the parent node, or None if this is the root level (default is None). -# start_id : int, optional -# The next available integer ID to assign to a new node (default is 0). -# include_failed_pp : bool, optional -# If False, nodes whose pretty-print text is "" are removed -# (default is True). -# deduplicate : bool, optional -# If True, deduplicate chains of identical nodes based on goalsBefore, -# goalsAfter, and pp (default is False). - -# Returns -# ------- -# nodes : dict[int, dict] -# A dictionary of node_id -> node_content. -# edges : list[tuple[int, int, dict]] -# A list of tuples (parent_id, child_id, {}) representing edges in the infotree. -# next_id : int -# The next available integer ID after processing all children. -# """ -# nodes = {} -# edges = [] - -# current_id = start_id - -# for item in infotree: -# if "node" in item: -# node_data = item["node"] - -# # Add this node to the nodes dictionary -# node_id = current_id -# current_id += 1 -# nodes[node_id] = node_data - -# # Recursively handle the children of the node, depth-first approach -# if ( -# "children" in item -# and isinstance(item["children"], list) -# and item["children"] -# ): -# child_nodes, child_edges, current_id = extract_nodes_and_edges( -# item["children"], -# parent_id=node_id, -# start_id=current_id, -# include_failed_pp=include_failed_pp, -# deduplicate=deduplicate, -# ) -# nodes.update(child_nodes) -# edges.extend(child_edges) - -# # Add an edge from the parent to this node -# if parent_id is not None: -# edges.append((parent_id, node_id, {})) - -# # Now handle possible flattening in a loop until there are no more changes -# transformed = True -# while transformed: -# transformed = False - -# # 1) Remove/flatten all children that have failed PP (if include_failed_pp=False) -# # We handle them individually, even if there's more than one child. -# # Then we break to re-check children from scratch, as new failed PP children might appear. -# child_list = [ -# e for e in edges if e[0] == node_id -# ] # edges from node_id to child -# for edge_obj in child_list: -# child_id = edge_obj[1] -# child_content = nodes.get(child_id, {}) -# child_pp = child_content.get("stx", {}).get("pp", "") - -# if not include_failed_pp and child_pp == "": -# # Flatten this child: remove it and connect parent directly to its children -# # Note that the new children might also have failed PP, so we need to re-check them -# _flatten_chain(nodes, edges, node_id, child_id) -# transformed = True -# break - -# if transformed: -# # We need to restart the while loop to re-check children from scratch, as there might be -# # more failed PP children -# continue - -# # 2) Deduplicate if there's exactly one child left that has same (goalsBefore, goalsAfter, pp) -# # This happens quite often in infotrees, where a tactic is repeated multiple times, extracted -# # with different parsers -# child_list = [e for e in edges if e[0] == node_id] -# if deduplicate and len(child_list) == 1: -# child_id = child_list[0][1] -# if child_id in nodes: -# child_content = nodes[child_id] -# # Compare parent's vs child's (goalsBefore, goalsAfter, pp) -# parent_goalsBefore = node_data.get("goalsBefore", []) -# parent_goalsAfter = node_data.get("goalsAfter", []) -# parent_pp = node_data.get("stx", {}).get("pp", "") - -# child_goalsBefore = child_content.get("goalsBefore", []) -# child_goalsAfter = child_content.get("goalsAfter", []) -# child_pp = child_content.get("stx", {}).get("pp", "") - -# if ( -# child_goalsBefore == parent_goalsBefore -# and child_goalsAfter == parent_goalsAfter -# and child_pp == parent_pp -# ): -# # Flatten this child: remove it and connect parent directly to its children -# _flatten_chain(nodes, edges, node_id, child_id) -# transformed = True - -# else: -# # If the item does not contain a 'node' key but might have children -# if "children" in item and isinstance(item["children"], list): -# child_nodes, child_edges, current_id = extract_nodes_and_edges( -# item["children"], -# parent_id=parent_id, -# start_id=current_id, -# include_failed_pp=include_failed_pp, -# deduplicate=deduplicate, -# ) -# nodes.update(child_nodes) -# edges.extend(child_edges) - -# return nodes, edges, current_id - - -# def _flatten_chain(nodes: dict, edges: list[tuple], parent_id: int, child_id: int): -# """ -# Flatten a chain by removing 'child_id' node and connecting 'parent_id' -# directly to the child's children. -# Given a parent node that has a single child, this function removes the child -# node from the dictionary of nodes and reassigns the child's children to the parent. -# This is used for node deduplication and removing failed-pp nodes. - -# Parameters -# ---------- -# nodes : dict[int, dict] -# A dictionary of node_id -> node_content. -# edges : list[tuple[int, int, dict]] -# A list of tuples (parent_id, child_id, {}) representing edges in the infotree. -# parent_id : int -# The ID of the parent node. -# child_id : int -# The ID of the child node that should be removed. - -# Returns -# ------- -# None -# This function modifies the nodes and edges in place. -# """ -# if child_id not in nodes: -# return - -# # Remove the node from the dictionary -# del nodes[child_id] - -# # Remove edge from parent_id -> child_id -# edges[:] = [e for e in edges if not (e[0] == parent_id and e[1] == child_id)] - -# # Reassign child's children edges to the parent -# for i, (src, tgt, attr) in enumerate(edges): -# if src == child_id: -# edges[i] = (parent_id, tgt, attr) - - -# def get_intervals(nodes: dict) -> list[dict]: -# """ -# Build a list of intervals from a given nodes dictionary. -# Each interval represents a tactic in the Lean file, capturing its -# start and finish positions, as well as the associated goals. - -# Parameters -# ---------- -# nodes : dict of {int : dict} -# A dictionary of node_id -> node_content. - -# Returns -# ------- -# intervals : list of dict -# A list of dictionaries, each containing: -# node_id, pp, start_line, start_col, finish_line, finish_col, goalsBefore, goalsAfter -# """ -# intervals = [] -# for node_id, node_content in nodes.items(): -# stx_range = node_content.get("stx", {}).get("range", {}) -# start_dict = stx_range.get("start", {}) -# finish_dict = stx_range.get("finish", {}) - -# intervals.append( -# { -# "node_id": node_id, -# "pp": node_content.get("stx", {}).get("pp", ""), -# "start_line": start_dict.get("line", 0), -# "start_col": start_dict.get("column", 0), -# "finish_line": finish_dict.get("line", 0), -# "finish_col": finish_dict.get("column", 0), -# "goalsBefore": node_content.get("goalsBefore", []), -# "goalsAfter": node_content.get("goalsAfter", []), -# } -# ) -# return intervals - - -# def adjust_intervals(intervals: list[dict]) -> list[dict]: -# """ -# Make intervals disjoint and create a file partition. -# Sort intervals by starting position, then set each interval's end to the next -# interval's start. This creates a sequence of adjacent intervals covering the file. - -# Parameters -# ---------- -# intervals : list of dict -# A list of dictionaries, each containing: -# node_id, pp, start_line, start_col, finish_line, finish_col, goalsBefore, goalsAfter - -# Returns -# ------- -# intervals_sorted : list of dict -# The updated intervals, sorted and trimmed so that they do not overlap. -# """ -# intervals_sorted = sorted( -# intervals, key=lambda iv: (iv["start_line"], iv["start_col"]) -# ) - -# # Remember the furthest original finish position -# max_finish_line, max_finish_col = -1, -1 -# for iv in intervals_sorted: -# if (iv["finish_line"], iv["finish_col"]) > (max_finish_line, max_finish_col): -# max_finish_line, max_finish_col = iv["finish_line"], iv["finish_col"] - -# for i in range(len(intervals_sorted) - 1): -# current = intervals_sorted[i] -# nxt = intervals_sorted[i + 1] -# current["finish_line"] = nxt["start_line"] -# current["finish_col"] = nxt["start_col"] -# current["goalsAfter"] = nxt["goalsBefore"] - -# if intervals_sorted: -# intervals_sorted[-1]["finish_line"] = max_finish_line -# intervals_sorted[-1]["finish_col"] = max_finish_col - -# intervals_sorted = [ -# iv -# for iv in intervals_sorted -# if not ( -# iv["start_line"] == iv["finish_line"] -# and iv["start_col"] == iv["finish_col"] -# ) -# ] - -# return intervals_sorted - - -# def retrieve_tactics(intervals: list[dict], source_lines: list[str]) -> list[dict]: -# """ -# Extract tactic code snippets from source lines based on intervals. - -# Parameters -# ---------- -# intervals : list of dict -# A list of dictionaries, each containing: -# node_id, pp, start_line, start_col, finish_line, finish_col, goalsBefore, goalsAfter -# Note: At this point, the pp field does not exactly correspond to the positions. -# source_lines : list of str -# The lines of the Lean file, read into a list. - -# Returns -# ------- -# results : list of dict -# A list of intervals augmented with the 'tactic' text from the file. -# Each dict has keys: goalsBefore, goalsAfter, tactic. -# """ -# results = [] -# for i in range(len(intervals)): -# iv = intervals[i] -# snippet_text = _extract_snippet( -# source_lines, -# iv["start_line"], -# iv["start_col"], -# iv["finish_line"], -# iv["finish_col"], -# ) -# data = { -# "goalsBefore": iv["goalsBefore"], -# "goalsAfter": iv["goalsAfter"], -# "tactic": snippet_text, -# } -# results.append(data) - -# return results - - -# def _extract_snippet( -# source_lines: list[str], -# start_line: int, -# start_col: int, -# finish_line: int, -# finish_col: int, -# ) -> str: -# """ -# Extract a code snippet from the Lean source lines. - -# Given a start and finish line-column pair, slice the lines to produce the exact text -# range in the Lean file. This handles both single-line and multi-line cases. - -# Parameters -# ---------- -# source_lines : list of str -# The lines read from the Lean file. -# start_line : int -# The 1-based starting line index. -# start_col : int -# The 0-based starting column index within start_line. -# finish_line : int -# The 1-based finishing line index. -# finish_col : int -# The 0-based finishing column index within finish_line. - -# Returns -# ------- -# str -# The extracted snippet from the Lean file, spanning (start_line, start_col) -# to (finish_line, finish_col). -# """ -# # Single line case -# if start_line == finish_line: -# line_idx = start_line - 1 -# line_text = source_lines[line_idx] -# return line_text[start_col:finish_col] - -# # Multi-line case -# # 1) from start_col to end-of-line for start_line -# snippet_parts = [] -# start_line_idx = start_line - 1 -# line_text = source_lines[start_line_idx] -# snippet_parts.append(line_text[start_col:]) - -# # 2) full lines between (start_line+1) .. (finish_line-1) -# for line_idx in range(start_line_idx + 1, finish_line - 1): -# snippet_parts.append(source_lines[line_idx]) - -# # 3) from begin-of-line up to finish_col for finish_line -# last_line_idx = finish_line - 1 -# last_line = source_lines[last_line_idx] -# snippet_parts.append(last_line[:finish_col]) - -# return "".join(snippet_parts) - - -# def separate_trailing_whitespace(s: str) -> tuple[str, str]: -# """ -# Remove trailing whitespace from a tactic and return (code, trailing_ws). - -# Example: -# - Input: "have h1 : ... := by\n " -# - Output: ("have h1 : ... := by", "\n ") - -# Parameters -# ---------- -# s : str -# A tactic string (Lean4 code). - -# Returns -# ------- -# code : str -# The tactic string without potential whitespaces at the end. -# trailing_ws : str -# The whitespaces at the end of s. Potentially None. -# """ -# code = s.rstrip(" \t\n\r") -# trailing_ws = s[len(code) :] -# return code, trailing_ws - - -# def separate_trailing_comment(s: str) -> tuple[str, str]: -# """ -# Remove trailing comments only if they start at the beginning of the line -# (modulo whitespace). Return (clean, trailing_comment). - -# A trailing comment can be: -# - One or more single-line comments (lines starting with '--'). -# - One multi-line block comment starting with '/-' (at line start, ignoring whitespace) -# and ending somewhere before the end of the string (must contain '-/'). - -# Example: -# - Input: "have h1 : ... := by\n -- To prove this have statement, we will just apply mul_pos" -# - Output: ("have h1 : ... := by\n ", "-- To prove this have statement, we will just apply mul_pos") - -# Note that the function does not extract trailing whitespaces from the tactic right before the start of the comment. - -# Parameters -# ---------- -# s : str -# A tactic string (Lean4 code). - -# Returns -# ------- -# clean : str -# The tactic string without potential comments at the end. -# comment : str -# The comment at the end of s. Potentially None. -# """ -# lines = s.splitlines(keepends=True) - -# # If no lines, do nothing -# if not lines: -# return s, "" - -# trailing_comments = [] # This will contain all the trailing comments -# i = len(lines) - 1 # Start from the end of the file - -# while i >= 0: -# line = lines[i] - -# # 1. If the line is blank, add it to the trailing_comments and go to the next line -# if line.strip() == "": -# trailing_comments.append(line) -# i -= 1 -# continue - -# # 2. Check if the line ends a multi-line block comment (i.e., has '-/'): -# if "-/" in line: -# end_idx = i -# start_idx = end_idx -# # Move upward until we find a line that starts with '/-' (mod whitespace) -# while start_idx >= 0 and not re.match(r"^\s*/-", lines[start_idx]): -# start_idx -= 1 -# if start_idx < 0: -# # Found '-/' but no valid start => not truly trailing -# break -# # Everything from start_idx..end_idx is the trailing block comment -# block_str = "".join(lines[start_idx : end_idx + 1]) -# trailing_comments.append(block_str) -# i = start_idx - 1 # jump above the block -# continue - -# # 3. Check if the line is a single-line comment -# if re.match(r"^\s*--", line) is not None: -# trailing_comments.append(line) -# i -= 1 -# continue - -# # 4. If none of the above, then we’ve hit a real line of code -# break - -# clean_code = "".join(lines[: i + 1]) - -# # If we didn't find any trailing comment, return the code with an empty string -# if trailing_comments == []: -# return clean_code, "" - -# else: -# trailing_comments.reverse() # trailing_comments is in bottom-to-top order, reverse it to restore top-to-bottom -# trailing_comments = "".join(trailing_comments) -# return clean_code, trailing_comments - - -# def transfer_trailing_whitespaces_and_comments(intervals: list[dict]): -# """ -# For each tactic in 'intervals' (except the last), remove trailing whitespaces, -# then trailing comments (if any), then trailing whitespaces again, -# and prepend all that to the next interval. - -# Parameters -# ---------- -# intervals : list of dict -# A list of dictionaries, each containing: -# goalsBefore, goalsAfter, tactic. - -# Returns -# ------- -# None -# The 'tactic' fields in intervals are modified in place. -# """ -# for i in range(len(intervals) - 1): -# current = intervals[i] -# nxt = intervals[i + 1] - -# # 1) Remove trailing whitespace -# code, trailing_ws_1 = separate_trailing_whitespace(current["tactic"]) - -# # 2) Remove trailing comment -# code, trailing_comment = separate_trailing_comment(code) - -# # 3) Remove trailing whitespace again -# code, trailing_ws_2 = separate_trailing_whitespace(code) - -# # Update the current tactic -# current["tactic"] = code - -# # Prepend to the next interval -# nxt["tactic"] = trailing_ws_2 + trailing_comment + trailing_ws_1 + nxt["tactic"] - - -# def remove_lean_comments(text) -> str: -# """ -# Remove single-line and multi-line comments from `text`. - -# Parameters -# ---------- -# text : str -# The Lean code that may contain comments. - -# Returns -# ------- -# text : str -# The Lean code without comments. -# """ -# # First, remove all multi-line comments -# pattern = re.compile(r"/-.*?-/", re.DOTALL) -# text = pattern.sub("", text) - -# # Then, remove all single-line comments -# lines = text.splitlines() -# new_lines = [] -# for line in lines: -# if line.strip() == "": -# new_lines.append(line) -# continue -# if "--" in line: -# # Keep only the part before the first occurrence of "--" -# line = line.split("--")[0].rstrip() -# if line.strip() != "": -# new_lines.append(line) -# text = "\n".join(new_lines) - -# return text - - -# def is_balanced(tactic: str) -> bool: -# """ -# Check whether `[]`, `()` and `⟨⟩` are balanced in `tactic`. - -# Parameters -# ---------- -# tactic : str -# A Lean tactic snippet. - -# Returns -# ------- -# bool -# True if every opening bracket/parenthesis is matched by its closing counterpart, False otherwise. -# """ -# return ( -# tactic.count("[") == tactic.count("]") -# and (tactic.count("(") == tactic.count(")")) -# and (tactic.count("⟨") == tactic.count("⟩")) -# ) - - -# def is_by(tactic: str) -> bool: -# """ -# Check whether the tactic is exactly the keyword `by`. -# Comments are removed and surrounding whitespace is ignored before the comparison. - -# Parameters -# ---------- -# tactic : str -# A Lean tactic snippet. - -# Returns -# ------- -# bool -# True if the tactic is a 'by', False otherwise. -# """ -# return remove_lean_comments(tactic).strip() == "by" - - -# def is_calc(tactic: str) -> bool: -# """ -# Check whether the snippet introduces a `calc` block. -# Accepted forms are either `calc` or `by calc`. - -# Parameters -# ---------- -# tactic : str -# A Lean tactic snippet. - -# Returns -# ------- -# bool -# True if the tactic introduces a `calc` block, False otherwise. -# """ -# s = remove_lean_comments(tactic).strip() -# # Direct match -# if s == "calc": -# return True -# # Match `by calc` with spaces or newlines after `by` -# m = re.match(r"^by\s+([\w_]+)$", s) -# return bool(m and m.group(1) == "calc") - - -# WRAPPER_TACTICS: set[str] = { -# "all_goals", -# "any_goals", -# "repeat", -# } - - -# def is_wrapper(tactic: str) -> bool: -# """ -# Check whether the tactic is a wrapper tactic (e.g., `all_goals`, `any_goals`, `repeat`), -# optionally preceded by `by`. - -# Parameters -# ---------- -# tactic : str -# A Lean tactic snippet. - -# Returns -# ------- -# bool -# True if the tactic is a wrapper tactic, False otherwise. -# """ -# s = remove_lean_comments(tactic).strip() -# # Direct match -# if s in WRAPPER_TACTICS: -# return True -# # Match `by ` with spaces or newlines after `by` -# m = re.match(r"^by\s+([\w_]+)$", s) -# return bool(m and m.group(1) in WRAPPER_TACTICS) - - -# def ends_with_by(tactic: str) -> bool: -# """ -# Check whether the tactic ends with `:= by` or `:=by`, ignoring trailing comments and whitespace. - -# Parameters -# ---------- -# tactic : str -# A Lean tactic snippet. - -# Returns -# ------- -# bool -# True if the tactic ends with `:= by` or `:=by`, False otherwise. -# """ -# s = remove_lean_comments(tactic).rstrip() -# return s.endswith(":=by") or s.endswith(":= by") - - -# def merge_intervals(intervals: list[dict]) -> list[dict]: -# """ -# Merge intervals that are not balanced or contain specific tactics. - -# Parameters -# ---------- -# intervals : list of dict -# A list of dictionaries, each containing: -# goalsBefore, goalsAfter, tactic. - -# Returns -# ------- -# merged_intervals : list of dict -# A list of dictionaries, each containing: -# goalsBefore, goalsAfter, tactic. -# The tactics that do not change the goals are merged with their successor. -# """ -# merged_intervals = [] -# i = 0 -# while i < len(intervals): -# accumulated = intervals[i]["tactic"] -# j = i + 1 -# # Merge subsequent intervals until accumulated tactic is balanced. -# while j < len(intervals) and ( -# not is_balanced(accumulated) -# or is_by(accumulated) -# or is_calc(accumulated) -# or is_wrapper(accumulated) -# ): -# accumulated += intervals[j]["tactic"] -# j += 1 - -# merged_interval = { -# "goalsBefore": intervals[i]["goalsBefore"], -# "goalsAfter": intervals[j - 1]["goalsAfter"], -# "tactic": accumulated, -# } - -# merged_intervals.append(merged_interval) -# i = j - -# # Transfer trailing `by` -# for k in range(len(merged_intervals) - 1): -# if ends_with_by(merged_intervals[k]["tactic"]): -# txt = merged_intervals[k]["tactic"] -# cut = txt.rfind(":=") + 2 # keep the ':=' -# by_part = txt[cut:] # 'by' or ' by', incl. space -# merged_intervals[k]["tactic"] = txt[:cut] -# merged_intervals[k + 1]["tactic"] = ( -# by_part + merged_intervals[k + 1]["tactic"] -# ) - -# # Set goalsBefore of tactic n to be the first non-empty goalsBefore starting at n. -# for i in range(len(merged_intervals) - 1): -# new_goals = [] -# for j in range(i, len(merged_intervals)): -# if merged_intervals[j]["goalsBefore"] != []: -# new_goals = merged_intervals[j]["goalsBefore"] -# break -# merged_intervals[i]["goalsBefore"] = new_goals - -# # Set goalsAfter of tactic n to be the goalsBefore of tactic n+1. -# for i in range(len(merged_intervals) - 1): -# new_goals = [] -# for j in range(i + 1, len(merged_intervals)): -# if merged_intervals[j]["goalsBefore"] != []: -# new_goals = merged_intervals[j]["goalsBefore"] -# break -# merged_intervals[i]["goalsAfter"] = new_goals - -# return merged_intervals - - -# def extract_data(infotree: list[dict], source_code: str) -> list[dict]: -# """ -# Performs the whole extraction process from an infotree and the corresponding Lean4 code. - -# This function: -# - Extracts nodes and edges from the infotree, -# - Removes synthetic nodes, -# - Builds and adjusts intervals to partition the Lean4 code, -# - Retrieves tactics from the Lean4 code, -# - Transfers trailing whitespaces and comments between consecutive tactics, -# - Merge intervals that are not balanced or contain specific tactics. - -# Parameters -# ---------- -# infotree : list of dict -# A list of dictionaries each containing 'node' and optionally 'children'. -# source_code : str -# The Lean4 code for retrieving text snippets. - -# Returns -# ------- -# intervals : list of dict -# A list of dictionaries, each containing: -# goalsBefore, goalsAfter, tactic. -# The intervals are partitioned and cover the whole file. -# """ -# # 1. Extract nodes and edges -# nodes, _, _ = extract_nodes_and_edges( -# infotree, include_failed_pp=False, deduplicate=True -# ) - -# # 2. Filter out the synthetic nodes -# nodes = { -# k: v -# for k, v in nodes.items() -# if not v.get("stx", {}).get("range", {}).get("synthetic", False) -# } - -# # 3. Build raw intervals from nodes -# intervals = get_intervals(nodes) - -# # 4. Adjust intervals so they become disjoint and partition the proof -# intervals = adjust_intervals(intervals) - -# # 5. Load lines from the Lean file -# source_lines = source_code.split("\n") -# source_lines = [line + "\n" for line in source_lines[:-1]] + [source_lines[-1]] - -# # 6. Extract the tactic for each final interval -# intervals = retrieve_tactics(intervals, source_lines) - -# # 7. Transfer trailing whitespaces and comments -# transfer_trailing_whitespaces_and_comments(intervals) - -# # 8. Merge intervals that are not balanced or contain specific tactics -# intervals = merge_intervals(intervals) - -# return intervals +import re +from typing import Optional, Dict, List, Tuple, Any, TypedDict, TypeAlias +from server.split import split_snippet +# from kimina_client.models import Infotree + +class Position(TypedDict): + line: int + column: int + +class Range(TypedDict): + synthetic: bool + start: Position + finish: Position + +class Syntax(TypedDict): + range: Range + pp: str + +class NodeInfo(TypedDict, total=False): + stx: Syntax + name: str | None + goalsBefore: List[str] + goalsAfter: List[str] + +class InfotreeItem(TypedDict, total=False): + node: NodeInfo + kind: str + children: List["InfotreeItem"] + +Infotree: TypeAlias = List[InfotreeItem] + +def extract_nodes_and_edges( + infotree: Infotree, + parent_id: Optional[int] = None, + start_id: int = 0, + include_failed_pp: bool = True, + deduplicate: bool = False, +) -> Tuple[Dict[int, NodeInfo], + List[Tuple[int, int, Dict[str, Any]]], + int]: + """ + Recursively extract nodes and edges from an infotree. + + Parameters + ---------- + infotree : Infotree, + A list of InfotreeItems each containing 'node' and 'children'. + parent_id : Optional[int] + The ID of the parent node, or None if this is the root level (default is None). + start_id : int + The next available integer ID to assign to a new node (default is 0). + include_failed_pp : bool + If False, nodes whose pretty-print text is "" are removed + (default is True). + deduplicate : bool + If True, deduplicate chains of identical nodes based on goalsBefore, + goalsAfter, and pp (default is False). + + Returns + ------- + nodes : Dict[int, NodeInfo] + A dictionary of node_id -> node_content. + edges : List[Tuple[int, int, Dict[str, Any]]] + A list of tuples (parent_id, child_id, {}) representing edges in the infotree. + next_id : int + The next available integer ID after processing all children. + """ + nodes : Dict[int, NodeInfo] = {} + edges : List[Tuple[int, int, Dict[str, Any]]] = [] + + current_id = start_id + + for item in infotree: + if "node" in item: + node_data = item["node"] + + # Add this node to the nodes dictionary + node_id = current_id + current_id += 1 + nodes[node_id] = node_data + + # Recursively handle the children of the node, depth-first approach + children = item.get("children", []) + if children: + child_nodes, child_edges, current_id = extract_nodes_and_edges( + children, + parent_id=node_id, + start_id=current_id, + include_failed_pp=include_failed_pp, + deduplicate=deduplicate, + ) + nodes.update(child_nodes) + edges.extend(child_edges) + + # Add an edge from the parent to this node + if parent_id is not None: + edges.append((parent_id, node_id, {})) + + # Now handle possible flattening in a loop until there are no more changes + transformed = True + while transformed: + transformed = False + + # 1) Remove/flatten all children that have failed PP (if include_failed_pp=False) + # We handle them individually, even if there's more than one child. + # Then we break to re-check children from scratch, as new failed PP children might appear. + child_list = [ + e for e in edges if e[0] == node_id + ] # edges from node_id to child + for edge_obj in child_list: + child_id = edge_obj[1] + child_content = nodes.get(child_id, {}) + child_pp = child_content.get("stx", {}).get("pp", "") + + if not include_failed_pp and child_pp == "": + # Flatten this child: remove it and connect parent directly to its children + # Note that the new children might also have failed PP, so we need to re-check them + _flatten_chain(nodes, edges, node_id, child_id) + transformed = True + break + + if transformed: + # We need to restart the while loop to re-check children from scratch, as there might be + # more failed PP children + continue + + # 2) Deduplicate if there's exactly one child left that has same (goalsBefore, goalsAfter, pp) + # This happens quite often in infotrees, where a tactic is repeated multiple times, extracted + # with different parsers + child_list = [e for e in edges if e[0] == node_id] + if deduplicate and len(child_list) == 1: + child_id = child_list[0][1] + if child_id in nodes: + child_content = nodes[child_id] + # Compare parent's vs child's (goalsBefore, goalsAfter, pp) + parent_goalsBefore = node_data.get("goalsBefore", []) + parent_goalsAfter = node_data.get("goalsAfter", []) + parent_pp = node_data.get("stx", {}).get("pp", "") + + child_goalsBefore = child_content.get("goalsBefore", []) + child_goalsAfter = child_content.get("goalsAfter", []) + child_pp = child_content.get("stx", {}).get("pp", "") + + if ( + child_goalsBefore == parent_goalsBefore + and child_goalsAfter == parent_goalsAfter + and child_pp == parent_pp + ): + # Flatten this child: remove it and connect parent directly to its children + _flatten_chain(nodes, edges, node_id, child_id) + transformed = True + + else: + # If the item does not contain a 'node' key but might have children + children = item.get("children") + if children: + child_nodes, child_edges, current_id = extract_nodes_and_edges( + children, + parent_id=parent_id, + start_id=current_id, + include_failed_pp=include_failed_pp, + deduplicate=deduplicate, + ) + nodes.update(child_nodes) + edges.extend(child_edges) + + return nodes, edges, current_id + + +def _flatten_chain(nodes: Dict[int, NodeInfo], edges: List[Tuple[int, int, Dict[str, Any]]], parent_id: int, child_id: int) -> None: + """ + Flatten a chain by removing 'child_id' node and connecting 'parent_id' + directly to the child's children. + Given a parent node that has a single child, this function removes the child + node from the dictionary of nodes and reassigns the child's children to the parent. + This is used for node deduplication and removing failed-pp nodes. + + Parameters + ---------- + nodes : Dict[int, NodeInfo] + A dictionary of node_id -> node_content. + edges : List[Tuple[int, int, Dict[str, Any]]] + A list of tuples (parent_id, child_id, {}) representing edges in the infotree. + parent_id : int + The ID of the parent node. + child_id : int + The ID of the child node that should be removed. + + Returns + ------- + None + This function modifies the nodes and edges in place. + """ + if child_id not in nodes: + return + + # Remove the node from the dictionary + del nodes[child_id] + + # Remove edge from parent_id -> child_id + edges[:] = [e for e in edges if not (e[0] == parent_id and e[1] == child_id)] + + # Reassign child's children edges to the parent + for i, (src, tgt, attr) in enumerate(edges): + if src == child_id: + edges[i] = (parent_id, tgt, attr) + + +def get_intervals(nodes: Dict[int, NodeInfo]) -> List[Dict[str, Any]]: + """ + Build a list of intervals from a given nodes dictionary. + Each interval represents a tactic in the Lean file, capturing its + start and finish positions, as well as the associated goals. + + Parameters + ---------- + nodes : Dict[int, NodeInfo] + A dictionary of node_id -> node_content. + + Returns + ------- + intervals : List[Dict[str, Any]] + A list of dictionaries, each containing: + node_id, pp, start_line, start_col, finish_line, finish_col, goalsBefore, goalsAfter + """ + intervals : List[Dict[str, Any]] = [] + for node_id, node_content in nodes.items(): + stx_range = node_content.get("stx", {}).get("range", {}) + start_dict = stx_range.get("start", {}) + finish_dict = stx_range.get("finish", {}) + + intervals.append( + { + "node_id": node_id, + "pp": node_content.get("stx", {}).get("pp", ""), + "start_line": start_dict.get("line", 0), + "start_col": start_dict.get("column", 0), + "finish_line": finish_dict.get("line", 0), + "finish_col": finish_dict.get("column", 0), + "goalsBefore": node_content.get("goalsBefore", []), + "goalsAfter": node_content.get("goalsAfter", []), + } + ) + return intervals + + +def adjust_intervals(intervals: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Make intervals disjoint and create a file partition. + Sort intervals by starting position, then set each interval's end to the next + interval's start. This creates a sequence of adjacent intervals covering the file. + + Parameters + ---------- + intervals : List[Dict[str, Any]] + A list of dictionaries, each containing: + node_id, pp, start_line, start_col, finish_line, finish_col, goalsBefore, goalsAfter + + Returns + ------- + intervals_sorted : List[Dict[str, Any]] + The updated intervals, sorted and trimmed so that they do not overlap. + """ + intervals_sorted = sorted( + intervals, key=lambda iv: (iv["start_line"], iv["start_col"]) + ) + + # Remember the furthest original finish position + max_finish_line, max_finish_col = -1, -1 + for iv in intervals_sorted: + if (iv["finish_line"], iv["finish_col"]) > (max_finish_line, max_finish_col): + max_finish_line, max_finish_col = iv["finish_line"], iv["finish_col"] + + for i in range(len(intervals_sorted) - 1): + current = intervals_sorted[i] + nxt = intervals_sorted[i + 1] + current["finish_line"] = nxt["start_line"] + current["finish_col"] = nxt["start_col"] + current["goalsAfter"] = nxt["goalsBefore"] + + if intervals_sorted: + intervals_sorted[-1]["finish_line"] = max_finish_line + intervals_sorted[-1]["finish_col"] = max_finish_col + + intervals_sorted = [ + iv + for iv in intervals_sorted + if not ( + iv["start_line"] == iv["finish_line"] + and iv["start_col"] == iv["finish_col"] + ) + ] + + return intervals_sorted + + +def retrieve_tactics(intervals: List[Dict[str, Any]], source_lines: list[str]) -> list[Dict[str, Any]]: + """ + Extract tactic code snippets from source lines based on intervals. + + Parameters + ---------- + intervals : List[Dict[str, Any]] + A list of dictionaries, each containing: + node_id, pp, start_line, start_col, finish_line, finish_col, goalsBefore, goalsAfter + Note: At this point, the pp field does not exactly correspond to the positions. + source_lines : List[str] + The lines of the Lean file, read into a list. + + Returns + ------- + results : List[Dict[str, Any]] + A list of intervals augmented with the 'tactic' text from the file. + Each dict has keys: goalsBefore, goalsAfter, tactic. + """ + results : List[Dict[str, Any]] = [] + for i in range(len(intervals)): + iv = intervals[i] + snippet_text = _extract_snippet( + source_lines, + iv["start_line"], + iv["start_col"], + iv["finish_line"], + iv["finish_col"], + ) + data = { + "goalsBefore": iv["goalsBefore"], + "goalsAfter": iv["goalsAfter"], + "tactic": snippet_text, + } + results.append(data) + + return results + + +def _extract_snippet( + source_lines: List[str], + start_line: int, + start_col: int, + finish_line: int, + finish_col: int, +) -> str: + """ + Extract a code snippet from the Lean source lines. + + Given a start and finish line-column pair, slice the lines to produce the exact text + range in the Lean file. This handles both single-line and multi-line cases. + + Parameters + ---------- + source_lines : List[str] + The lines read from the Lean file. + start_line : int + The 1-based starting line index. + start_col : int + The 0-based starting column index within start_line. + finish_line : int + The 1-based finishing line index. + finish_col : int + The 0-based finishing column index within finish_line. + + Returns + ------- + str + The extracted snippet from the Lean file, spanning (start_line, start_col) + to (finish_line, finish_col). + """ + # Single line case + if start_line == finish_line: + line_idx = start_line - 1 + line_text = source_lines[line_idx] + return line_text[start_col:finish_col] + + # Multi-line case + # 1) from start_col to end-of-line for start_line + snippet_parts : List[str] = [] + start_line_idx = start_line - 1 + line_text = source_lines[start_line_idx] + snippet_parts.append(line_text[start_col:]) + + # 2) full lines between (start_line+1) .. (finish_line-1) + for line_idx in range(start_line_idx + 1, finish_line - 1): + snippet_parts.append(source_lines[line_idx]) + + # 3) from begin-of-line up to finish_col for finish_line + last_line_idx = finish_line - 1 + last_line = source_lines[last_line_idx] + snippet_parts.append(last_line[:finish_col]) + + return "".join(snippet_parts) + + +def separate_trailing_whitespace(s: str) -> Tuple[str, str]: + """ + Remove trailing whitespace from a tactic and return (code, trailing_ws). + + Example: + - Input: "have h1 : ... := by\n " + - Output: ("have h1 : ... := by", "\n ") + + Parameters + ---------- + s : str + A tactic string (Lean4 code). + + Returns + ------- + code : str + The tactic string without potential whitespaces at the end. + trailing_ws : str + The whitespaces at the end of s. Potentially None. + """ + code = s.rstrip(" \t\n\r") + trailing_ws = s[len(code) :] + return code, trailing_ws + + +def separate_trailing_comment(s: str) -> Tuple[str, str]: + """ + Remove trailing comments only if they start at the beginning of the line + (modulo whitespace). Return (clean, trailing_comment). + + A trailing comment can be: + - One or more single-line comments (lines starting with '--'). + - One multi-line block comment starting with '/-' (at line start, ignoring whitespace) + and ending somewhere before the end of the string (must contain '-/'). + + Example: + - Input: "have h1 : ... := by\n -- To prove this have statement, we will just apply mul_pos" + - Output: ("have h1 : ... := by\n ", "-- To prove this have statement, we will just apply mul_pos") + + Note that the function does not extract trailing whitespaces from the tactic right before the start of the comment. + + Parameters + ---------- + s : str + A tactic string (Lean4 code). + + Returns + ------- + clean : str + The tactic string without potential comments at the end. + comment : str + The comment at the end of s. Potentially None. + """ + lines = s.splitlines(keepends=True) + + # If no lines, do nothing + if not lines: + return s, "" + + trailing_comments : List[str] = [] # This will contain all the trailing comments + i = len(lines) - 1 # Start from the end of the file + + while i >= 0: + line = lines[i] + + # 1. If the line is blank, add it to the trailing_comments and go to the next line + if line.strip() == "": + trailing_comments.append(line) + i -= 1 + continue + + # 2. Check if the line ends a multi-line block comment (i.e., has '-/'): + if "-/" in line: + end_idx = i + start_idx = end_idx + # Move upward until we find a line that starts with '/-' (mod whitespace) + while start_idx >= 0 and not re.match(r"^\s*/-", lines[start_idx]): + start_idx -= 1 + if start_idx < 0: + # Found '-/' but no valid start => not truly trailing + break + # Everything from start_idx..end_idx is the trailing block comment + block_str = "".join(lines[start_idx : end_idx + 1]) + trailing_comments.append(block_str) + i = start_idx - 1 # jump above the block + continue + + # 3. Check if the line is a single-line comment + if re.match(r"^\s*--", line) is not None: + trailing_comments.append(line) + i -= 1 + continue + + # 4. If none of the above, then we’ve hit a real line of code + break + + clean_code = "".join(lines[: i + 1]) + + # If we didn't find any trailing comment, return the code with an empty string + if trailing_comments == []: + return clean_code, "" + + else: + trailing_comments.reverse() # trailing_comments is in bottom-to-top order, reverse it to restore top-to-bottom + comment = "".join(trailing_comments) + return clean_code, comment + + +def transfer_trailing_whitespaces_and_comments(intervals: List[Dict[str, Any]]) -> None: + """ + For each tactic in 'intervals' (except the last), remove trailing whitespaces, + then trailing comments (if any), then trailing whitespaces again, + and prepend all that to the next interval. + + Parameters + ---------- + intervals : List[Dict[str, Any]] + A list of dictionaries, each containing: + goalsBefore, goalsAfter, tactic. + + Returns + ------- + None + The 'tactic' fields in intervals are modified in place. + """ + for i in range(len(intervals) - 1): + current = intervals[i] + nxt = intervals[i + 1] + + # 1) Remove trailing whitespace + code, trailing_ws_1 = separate_trailing_whitespace(current["tactic"]) + + # 2) Remove trailing comment + code, trailing_comment = separate_trailing_comment(code) + + # 3) Remove trailing whitespace again + code, trailing_ws_2 = separate_trailing_whitespace(code) + + # Update the current tactic + current["tactic"] = code + + # Prepend to the next interval + nxt["tactic"] = trailing_ws_2 + trailing_comment + trailing_ws_1 + nxt["tactic"] + + +def remove_lean_comments(text: str) -> str: + """ + Remove single-line and multi-line comments from `text`. + + Parameters + ---------- + text : str + The Lean code that may contain comments. + + Returns + ------- + text : str + The Lean code without comments. + """ + # First, remove all multi-line comments + pattern = re.compile(r"/-.*?-/", re.DOTALL) + text = pattern.sub("", text) + + # Then, remove all single-line comments + lines = text.splitlines() + new_lines : List[str] = [] + for line in lines: + if line.strip() == "": + new_lines.append(line) + continue + if "--" in line: + # Keep only the part before the first occurrence of "--" + line = line.split("--")[0].rstrip() + if line.strip() != "": + new_lines.append(line) + text = "\n".join(new_lines) + + return text + + +def is_balanced(tactic: str) -> bool: + """ + Check whether `[]`, `()` and `⟨⟩` are balanced in `tactic`. + + Parameters + ---------- + tactic : str + A Lean tactic snippet. + + Returns + ------- + bool + True if every opening bracket/parenthesis is matched by its closing counterpart, False otherwise. + """ + return ( + tactic.count("[") == tactic.count("]") + and (tactic.count("(") == tactic.count(")")) + and (tactic.count("⟨") == tactic.count("⟩")) + ) + + +def is_by(tactic: str) -> bool: + """ + Check whether the tactic is exactly the keyword `by`. + Comments are removed and surrounding whitespace is ignored before the comparison. + + Parameters + ---------- + tactic : str + A Lean tactic snippet. + + Returns + ------- + bool + True if the tactic is a 'by', False otherwise. + """ + return remove_lean_comments(tactic).strip() == "by" + + +def is_calc(tactic: str) -> bool: + """ + Check whether the snippet introduces a `calc` block. + Accepted forms are either `calc` or `by calc`. + + Parameters + ---------- + tactic : str + A Lean tactic snippet. + + Returns + ------- + bool + True if the tactic introduces a `calc` block, False otherwise. + """ + s = remove_lean_comments(tactic).strip() + # Direct match + if s == "calc": + return True + # Match `by calc` with spaces or newlines after `by` + m = re.match(r"^by\s+([\w_]+)$", s) + return bool(m and m.group(1) == "calc") + + +WRAPPER_TACTICS: set[str] = { + "all_goals", + "any_goals", + "repeat", +} + + +def is_wrapper(tactic: str) -> bool: + """ + Check whether the tactic is a wrapper tactic (e.g., `all_goals`, `any_goals`, `repeat`), + optionally preceded by `by`. + + Parameters + ---------- + tactic : str + A Lean tactic snippet. + + Returns + ------- + bool + True if the tactic is a wrapper tactic, False otherwise. + """ + s = remove_lean_comments(tactic).strip() + # Direct match + if s in WRAPPER_TACTICS: + return True + # Match `by ` with spaces or newlines after `by` + m = re.match(r"^by\s+([\w_]+)$", s) + return bool(m and m.group(1) in WRAPPER_TACTICS) + + +def ends_with_by(tactic: str) -> bool: + """ + Check whether the tactic ends with `:= by` or `:=by`, ignoring trailing comments and whitespace. + + Parameters + ---------- + tactic : str + A Lean tactic snippet. + + Returns + ------- + bool + True if the tactic ends with `:= by` or `:=by`, False otherwise. + """ + s = remove_lean_comments(tactic).rstrip() + return s.endswith(":=by") or s.endswith(":= by") + + +def merge_intervals(intervals: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Merge intervals that are not balanced or contain specific tactics. + + Parameters + ---------- + intervals : List[Dict[str, Any]] + A list of dictionaries, each containing: + goalsBefore, goalsAfter, tactic. + + Returns + ------- + merged_intervals : List[Dict[str, Any]] + A list of dictionaries, each containing: + goalsBefore, goalsAfter, tactic. + The tactics that do not change the goals are merged with their successor. + """ + merged_intervals : List[Dict[str, Any]] = [] + i = 0 + while i < len(intervals): + accumulated = intervals[i]["tactic"] + j = i + 1 + # Merge subsequent intervals until accumulated tactic is balanced. + while j < len(intervals) and ( + not is_balanced(accumulated) + or is_by(accumulated) + or is_calc(accumulated) + or is_wrapper(accumulated) + ): + accumulated += intervals[j]["tactic"] + j += 1 + + merged_interval = { + "goalsBefore": intervals[i]["goalsBefore"], + "goalsAfter": intervals[j - 1]["goalsAfter"], + "tactic": accumulated, + } + + merged_intervals.append(merged_interval) + i = j + + # Transfer trailing `by` + for k in range(len(merged_intervals) - 1): + if ends_with_by(merged_intervals[k]["tactic"]): + txt = merged_intervals[k]["tactic"] + cut = txt.rfind(":=") + 2 # keep the ':=' + by_part = txt[cut:] # 'by' or ' by', incl. space + merged_intervals[k]["tactic"] = txt[:cut] + merged_intervals[k + 1]["tactic"] = ( + by_part + merged_intervals[k + 1]["tactic"] + ) + + # Set goalsBefore of tactic n to be the first non-empty goalsBefore starting at n. + for i in range(len(merged_intervals) - 1): + new_goals = [] + for j in range(i, len(merged_intervals)): + if merged_intervals[j]["goalsBefore"] != []: + new_goals = merged_intervals[j]["goalsBefore"] + break + merged_intervals[i]["goalsBefore"] = new_goals + + # Set goalsAfter of tactic n to be the goalsBefore of tactic n+1. + for i in range(len(merged_intervals) - 1): + new_goals = [] + for j in range(i + 1, len(merged_intervals)): + if merged_intervals[j]["goalsBefore"] != []: + new_goals = merged_intervals[j]["goalsBefore"] + break + merged_intervals[i]["goalsAfter"] = new_goals + + return merged_intervals + + +def extract_data(infotree: Infotree, source_code: str) -> List[Dict[str, Any]]: + """ + Performs the whole extraction process from an infotree and the corresponding Lean4 code. + + This function: + - Extracts nodes and edges from the infotree, + - Removes synthetic nodes, + - Builds and adjusts intervals to partition the Lean4 code, + - Retrieves tactics from the Lean4 code, + - Transfers trailing whitespaces and comments between consecutive tactics, + - Merge intervals that are not balanced or contain specific tactics. + + Parameters + ---------- + infotree : Infotree + A list of dictionaries each containing 'node' and optionally 'children'. + source_code : str + The Lean4 code for retrieving text snippets. + + Returns + ------- + intervals : List[Dict[str, Any]] + A list of dictionaries, each containing: + goalsBefore, goalsAfter, tactic. + The intervals are partitioned and cover the whole file. + """ + # 1. Extract nodes and edges + nodes, _, _ = extract_nodes_and_edges( + infotree, include_failed_pp=False, deduplicate=True + ) + + # 2. Filter out the synthetic nodes + nodes = { + k: v + for k, v in nodes.items() + if not v.get("stx", {}).get("range", {}).get("synthetic", False) + } + + # 3. Build raw intervals from nodes + intervals = get_intervals(nodes) + + # 4. Adjust intervals so they become disjoint and partition the proof + intervals = adjust_intervals(intervals) + + # 5. Load lines from the Lean file + source_code = split_snippet(source_code)[1] + source_lines = source_code.split("\n") + source_lines = [line + "\n" for line in source_lines[:-1]] + [source_lines[-1]] + + # 6. Extract the tactic for each final interval + intervals = retrieve_tactics(intervals, source_lines) + + # 7. Transfer trailing whitespaces and comments + transfer_trailing_whitespaces_and_comments(intervals) + + # 8. Merge intervals that are not balanced or contain specific tactics + intervals = merge_intervals(intervals) + + return intervals \ No newline at end of file diff --git a/client/kimina_client/utils.py b/client/kimina_client/utils.py index 0eccea3..2320beb 100644 --- a/client/kimina_client/utils.py +++ b/client/kimina_client/utils.py @@ -1,5 +1,6 @@ import logging from difflib import get_close_matches +from typing import Tuple from colorama import Style @@ -98,3 +99,46 @@ def build_log(dataset_name: str, n: int, batch_size: int) -> str: def b(s: str) -> str: return str(Style.BRIGHT + s + Style.RESET_ALL) + + +def split_snippet(code: str) -> Tuple[str, str]: + """ + Splits a code snippet into a header (imports) and body. + + - Header: all lines at the top that are 'import ...' or blank before the first non-import line. + If any import starts with 'import Mathlib', include a single 'import Mathlib' at the top of the header. + Other imports follow in their original order, without duplicates. + - Body: the rest of the code starting from the first non-import/non-blank line. + """ + lines = code.splitlines() + + # Separate header from body + i = 0 + while i < len(lines) and ( + lines[i].strip() == "" or lines[i].strip().startswith("import ") + ): + i += 1 + header_lines = [x.strip() for x in lines[:i]] + body = "\n".join(lines[i:]) + + # Process imports in header + import_lines = [line for line in header_lines if line.startswith("import ")] + imports: list[str] = [] + seen: set[str] = set() + has_mathlib = False + for line in import_lines: + if line.startswith("import Mathlib"): + has_mathlib = True + else: + if line not in seen: + seen.add(line) + imports.append(line) + + # Build final header + result_header: list[str] = [] + if has_mathlib: + result_header.append("import Mathlib") + result_header.extend(imports) + + header = "\n".join(result_header) + return header, body \ No newline at end of file diff --git a/setup.sh b/setup.sh index 26014ea..3f1f2bc 100755 --- a/setup.sh +++ b/setup.sh @@ -1,6 +1,12 @@ #!/usr/bin/env bash set -euxo pipefail +if [ -f .env ]; then + set -a + source .env + set +a +fi + LEAN_SERVER_LEAN_VERSION="${LEAN_SERVER_LEAN_VERSION:-v4.15.0}" REPL_REPO_URL="${REPL_REPO_URL:-https://github.com/leanprover-community/repl.git}" REPL_BRANCH="${REPL_BRANCH:-$LEAN_SERVER_LEAN_VERSION}" diff --git a/tests/test_infotree_processing.py b/tests/test_infotree_processing.py new file mode 100644 index 0000000..b429666 --- /dev/null +++ b/tests/test_infotree_processing.py @@ -0,0 +1,18 @@ +from kimina_client import KiminaClient +from kimina_client.infotree import extract_data + +def test_infotree_processing() -> None: + client = KiminaClient() + + with open('examples/demo_search_2.lean', 'r') as file: + code = file.read() + + check_response = client.check(code, infotree="original") + result = check_response.results[0] + infotree = result.response["infotree"] + data = extract_data(infotree, code) + + assert len(data) == 11 + assert data[0]["tactic"] == 'by\n have h₁ : 0 < a + b + c :=' + assert data[0]["goalsBefore"] == ['a b c : ℝ\nha : 0 < a\nhb : 0 < b\nhc : 0 < c\n⊢ a ^ 2 / (b + c) + b ^ 2 / (a + c) + 16 * c ^ 2 / (a + b) ≥ 1 / 9 * (64 * c - a - b)'] + assert data[0]["goalsAfter"] == ['a b c : ℝ\nha : 0 < a\nhb : 0 < b\nhc : 0 < c\n⊢ 0 < a + b + c'] \ No newline at end of file