upload files

prover/lean/proof.py

@@ -0,0 +1,183 @@
+import re
+import json
+
+import numpy as np
+
+from prover.utils import AttrDict, LEAN4_DEFAULT_HEADER
+
+
+class Proof(object):
+    def __init__(self, full_code, _args, _result_backup=None, **kwargs):
+        self._kwargs_backup = kwargs
+        for key, val in kwargs.items():
+            self.__setattr__(key, val)
+        self._args = _args
+        self._update_full_code(full_code, _result_backup=_result_backup)
+    
+    @property
+    def result(self):
+        if self._verifier_request_id is not None:
+            self._result = self._scheduler.verifier_get_request_outputs(self._verifier_request_id)
+            self._verifier_request_id = None
+        return self._result
+    
+    def is_result_ready(self):
+        if self._verifier_request_id is None:
+            return True
+        status = self._scheduler.verifier_get_request_status(self._verifier_request_id)
+        if status is not None:
+            self._result = status
+            self._verifier_request_id = None
+        return self._result is not None
+    
+    @property
+    def cleaned_code(self):
+        return self.full_code[len(self.header) + len(self.formal_statement): len(self.full_code) - len(self.tailer)]
+    
+    def _update_full_code(self, full_code, _result_backup=None):
+        self.full_code = full_code
+        self._verifier_request_id, self._result = None, None
+        if _result_backup is not None:
+            self._result = _result_backup
+        elif self._args.require_verification:  # need to call verification server
+            self._verifier_request_id = self._scheduler.verifier_submit_request(dict(
+                code=self.full_code,
+                ast=True, tactics=True,
+            ))
+        self._parse_full_code_lines()
+    
+    def _parse_full_code_lines(self):
+        self._full_code_lines = self.full_code.split('\n')
+        self._line_offset, _offset = [], -1
+        for _line in self._full_code_lines:
+            _offset += 1  # '\n'
+            self._line_offset.append(_offset)
+            _offset += len(_line)
+    
+    def _get_idx(self, pos_info):
+        return self._line_offset[pos_info['line'] - 1] + pos_info['column']
+    
+    def segmentation(self, result=None):
+        if result is None:
+            result = self.result
+        if 'errors' not in result:
+            # compiler timeout
+            return []
+        
+        _prefix_len = len(self.header) + len(self.formal_statement)
+        truncate_pos = len(self.full_code) - len(self.tailer)
+        for info in result['sorries'] + result['errors']:
+            info_pos = self._get_idx(info['pos'])
+            if info_pos >= _prefix_len and not info.get('data', str()).lstrip().startswith('unsolved goals'):
+                truncate_pos = min(truncate_pos, info_pos)
+        partial_code = self.full_code[:truncate_pos]
+        if len(partial_code) <= _prefix_len:
+            # all proof lines are invalid
+            return []
+        
+        code_lines = partial_code.split('\n')
+        pos_last, segments = _prefix_len, []
+        for line_idx in range(len(code_lines)):
+            if self._line_offset[line_idx] >= _prefix_len:
+                def compute_last_valid_char_pos(line):
+                    idx, last_non_blank = 0, len(line) + 1
+                    while idx < len(line):
+                        if line[idx: idx+2] == '--':
+                            return last_non_blank
+                        elif line[idx: idx+2] == '/-':
+                            if '-/' not in line[idx+2:]:
+                                # cannot split in this line
+                                return len(line) + 1
+                            idx = line.find('-/', idx+2) + 1
+                        elif line[idx] != ' ':
+                            last_non_blank = idx
+                        idx += 1
+                    return last_non_blank
+                line_lastChar = self._line_offset[line_idx] + compute_last_valid_char_pos(code_lines[line_idx])
+                line_endPos = self._line_offset[line_idx] + len(code_lines[line_idx])
+
+                pos_min, goal = 1e9, None
+                for tactic_info in result['ast']['tactics']:
+                    pos, endPos = tactic_info['pos'], tactic_info['endPos']
+                    if line_lastChar <= endPos and endPos <= line_endPos and pos < pos_min:
+                        pos_min = pos
+                        goal = tactic_info['stateAfter']
+                if goal is not None:
+                    for tactic_info in result['ast']['tactics']:
+                        pos, endPos = tactic_info['pos'], tactic_info['endPos']
+                        if pos_last < endPos and endPos <= line_endPos and pos < pos_min:
+                            pos_min = pos
+
+                    while pos_min > 0 and partial_code[pos_min - 1] != '\n':
+                        pos_min -= 1
+                    indent_len = 0
+                    while partial_code[pos_min + indent_len] == ' ':
+                        indent_len += 1
+                    newline_with_indent = '\n' + ' ' * indent_len
+                    
+                    segments.append(AttrDict(
+                        tactic_code=partial_code[pos_last: line_endPos] + '\n',
+                        state_comment=newline_with_indent.join([
+                            ' ' * indent_len + '/- tactic state:',
+                            '  ' + goal.replace('\n', newline_with_indent + '  '),
+                            '-/\n'
+                        ]),
+                        goal=goal,
+                        indent=indent_len,
+                    ))
+                    pos_last = line_endPos + 1
+        if result['complete'] and (len(segments) == 0 or segments[-1].goal != 'no goals' or segments[-1].indent != segments[0].indent):
+            indent_len = 2 if len(segments) == 0 else segments[0].indent
+            newline_with_indent = '\n' + ' ' * indent_len
+            segments.append(AttrDict(
+                tactic_code=partial_code[pos_last:].rstrip(' \n') + '\n',
+                state_comment=newline_with_indent.join([
+                    ' ' * indent_len + '/- tactic state:',
+                    '  no goals',
+                    '-/\n'
+                ]),
+                goal='no goals',
+                indent=indent_len,
+            ))
+        segments = [seg for seg in segments if len(seg.tactic_code.strip(' \n')) > 0]
+        return segments
+
+
+class ProofSummarizer(object):
+    def __init__(self, data, scheduler=None):
+        """
+        Inputs:
+            data (`dict`): The problem information storing in a `dict` object.
+                formal_statement (`str`): The formal statement of the unproved problem;
+                header (`str`, *optional*, defaults to ''): The code header required by the complier;
+                tailer (`str`, *optional*, defaults to ''): The code tailer required by the complier.
+            scheduler (prover.workers.scheduler.Scheduler, *optional*, defaults to None):
+                An interface to submit requests to models and the verification server.
+                If set to None, the downstream tasks may require the verification result as inputs.
+        """
+        self.formal_statement = data['formal_statement']
+        self.header = data.get('header', LEAN4_DEFAULT_HEADER)
+        self.tailer = data.get('tailer', str())
+        self.scheduler = scheduler
+    
+    def analyze(self, code, require_verification=True):
+        """
+        Inputs:
+            code (`str`): The code of formal proof.
+            require_verification (`bool`, *optional*, defaults to True):
+                Whether to submit a request to the verification server.
+                If set to False, the downstream tasks may require the verification result as inputs.
+        Return:
+            A `Proof` object that summarizes the code.
+        """
+        return Proof(
+            full_code=''.join([self.header, self.formal_statement, code.rstrip(' \n'), self.tailer]),
+            raw_code=code,
+            formal_statement=self.formal_statement,
+            header=self.header,
+            tailer=self.tailer,
+            _scheduler=self.scheduler,
+            _args=AttrDict(
+                require_verification=require_verification,
+            )
+        )

prover/lean/verifier.py

@@ -0,0 +1,156 @@
+import os
+import time
+import json
+import ctypes
+import resource
+import tempfile
+import traceback
+import threading
+import subprocess
+import multiprocessing as mp
+from pprint import pprint
+
+import numpy as np
+
+from prover.lean.ast_parser import lean4_parser
+from prover.workers import ProcessScheduler
+from prover.utils import AttrDict
+
+
+HOME_DIR = os.path.expanduser('~')
+DEFAULT_LAKE_PATH = f'{HOME_DIR}/.elan/bin/lake'
+DEFAULT_LEAN_WORKSPACE = 'mathlib4/'
+
+
+def verify_lean4_file(code, lake_path=DEFAULT_LAKE_PATH, lean_workspace=DEFAULT_LEAN_WORKSPACE, last_env=None, verbose=False, timeout=300, allTactics=False, ast=False, premises=False, tactics=False):
+    command = dict(cmd=code, allTactics=allTactics, ast=ast, tactics=tactics, premises=premises)
+    if last_env is not None:
+        command.update(env=last_env)
+    message_str = json.dumps(command, ensure_ascii=False)
+    if verbose:
+        print(message_str)
+    start_time = time.time()
+    system_messages = ''
+    try:
+        with tempfile.TemporaryFile(mode='w+', encoding='utf-8') as temp_file:
+            temp_file.write(message_str + "\r\n\r\n")
+            temp_file.seek(0)
+            outputs = subprocess.run([lake_path, "exe", 'repl'], stdin=temp_file, capture_output=True, text=True, cwd=lean_workspace, timeout=timeout)
+        result = json.loads(outputs.stdout)
+        ast_results = lean4_parser(code, result['ast']) if 'ast' in result and result['ast'] else {}
+        result = {
+            "sorries" : result.get('sorries', []), 
+            "tactics" : result.get('tactics', []),
+            "errors" : [m for m in result.get('messages', []) if m['severity'] == 'error'],
+            "warnings" : [m for m in result.get('messages', []) if m['severity'] == 'warning'],
+            "infos" : [m for m in result.get('messages', []) if m['severity'] == 'info'],
+            "system_messages" : system_messages,
+            "system_errors" : None,
+            "ast" : ast_results,
+            "verified_code" : code,
+        }
+        result['pass'] = not result['errors']
+        result['complete'] = result['pass'] and not result['sorries'] and not any("declaration uses 'sorry'" in warning['data'] or 'failed' in warning['data'] for warning in result['warnings'])
+    except:
+        result = {
+            "pass": False,
+            "complete": False,
+            "system_errors": traceback.format_exc(),
+            "system_messages": system_messages
+        }
+    result['verify_time'] = time.time() - start_time
+    return result
+
+
+class Lean4ServerProcess(mp.Process):
+    def __init__(self, idx, task_queue, request_statuses, lock, extra_args=AttrDict()):
+        super().__init__()
+        self.idx = idx
+        self.task_queue = task_queue
+        self.request_statuses = request_statuses
+        self.lock = lock
+        self.extra_args = extra_args
+
+        self.timeout = extra_args.get('timeout', 300)
+        self.memory_limit = extra_args.get('memory_limit', -1)
+        self.last_output_time = mp.Value(ctypes.c_double, time.time())
+        self.complete_count = mp.Value(ctypes.c_int, 0)
+    
+    def run(self):
+        if self.memory_limit > 0:
+            resource.setrlimit(
+                resource.RLIMIT_AS,
+                (self.memory_limit * (1000 ** 3), self.memory_limit * (1000 ** 3))
+            )
+        while True:
+            inputs = self.task_queue.get()
+            if inputs is None: # Terminate when receiving None
+                break
+            for _, request_id, task in inputs:
+                if isinstance(task, str):
+                    task = dict(code=task)
+                if 'timeout' not in task:
+                    task['timeout'] = self.timeout
+                result = verify_lean4_file(**task)
+                if len(result['system_messages']) > 0:
+                    retry_start_time = time.time()
+                    while ('lean::exception: failed to create thread' in result['system_messages'] or
+                           'std::bad_alloc: std::bad_alloc' in result['system_messages'] or
+                           'Cannot allocate memory' in result['system_messages']) \
+                          and time.time() - retry_start_time < self.timeout:
+                        time.sleep(0.1)
+                        result = verify_lean4_file(**task)
+                with self.lock:
+                    self.request_statuses[request_id] = result
+                    self.last_output_time.value = time.time()
+                    self.complete_count.value += 1
+
+
+class Lean4ServerScheduler(ProcessScheduler):
+    def __init__(self, max_concurrent_requests=64, timeout=300, memory_limit=-1, name='verifier'):
+        super().__init__(batch_size=1, name=name)
+        
+        self.processes = [
+            Lean4ServerProcess(
+                idx=idx,
+                task_queue=self.task_queue,
+                request_statuses=self.request_statuses,
+                lock=self.lock,
+                extra_args=AttrDict(
+                    timeout=timeout,
+                    memory_limit=memory_limit,
+                )
+            )
+            for idx in range(max_concurrent_requests)
+        ]
+        for p in self.processes:
+            p.start()
+        print(f'Complete launching {len(self.processes)} LeanServerProcesses')
+
+        self.timeout = timeout
+        self._running_monitor = mp.Value(ctypes.c_bool, True)
+        self._last_complete_count = mp.Value(ctypes.c_int, 0)
+        self._monitor_process = mp.Process(target=self._monitor)
+        self._monitor_process.start()
+    
+    def _monitor(self):
+        while self._running_monitor.value:
+            time.sleep(1.0)
+            subprocess.run(['killall', 'repl', f'--older-than={int(self.timeout) + 10}s'], capture_output=True)
+    
+    def close(self):
+        super().close()
+        for p in self.processes:
+            p.join()
+        self._running_monitor.value = False
+        self._monitor_process.join()
+        print(f'All {len(self.processes)} LeanServerProcesses stopped')
+
+
+if __name__ == '__main__':
+    code = open('mathlib4/.lake/packages/REPL/test/aime_1983_p9.code.in').read()
+    lean4_scheduler = Lean4ServerScheduler(max_concurrent_requests=1, timeout=300, memory_limit=10, name='verifier')
+    request_id_list = lean4_scheduler.submit_all_request([dict(code=code, ast=True, tactics=True)])
+    outputs_list = lean4_scheduler.get_all_request_outputs(request_id_list)
+    lean4_scheduler.close()
+    pprint(outputs_list)