1# NOTE: This file is derived from data_reader.pyx. Any changes here should be reflected there as well.
  2
  3"""This module provides read access to dataset files of a Tactican API dataset.
  4Additionally, there is support for communicating with a Coq processes through
  5the same data-format. The next sections give an overview of how this can be
  6accomplished. For an overview of examples, see `pytact`.
  7
  8# Reading a dataset
  9
 10The dataset is mapped into memory using mmap, allowing random access to it's
 11structures while keeping memory low. This file contains three entry-points to a
 12dataset in order of preference:
 13
 141. Contextmanager `pytact.data_reader.data_reader(path)` provides high-level access to the data
 15   in directory `path`. This is the preferred entry-point unless you need
 16   something special.
 172. Contextmanager `pytact.data_reader.lowlevel_data_reader(path)` provides low-level access to
 18   the data in directory `path`, giving direct access to the Cap'n Proto structures
 19   of the dataset. Use this when `data_reader` is too slow or you need access
 20   to data not provided by `data_reader`.
 213. Contextmanager `file_dataset_reader` provides low-level access to
 22   the Cap'n Proto structures of a single file. Using this is usually not
 23   advisable.
 24
 25Additionally, some indexing and traversal helpers are provided:
 26- `GlobalContextSets` calculates and caches the global context of a definition
 27  as a set, also caching intermediate results.
 28- `definition_dependencies` and `node_dependencies` traverse the graph starting
 29  from a node and return all direct, non-transitive definitions that node depends on.
 30
 31# Communicating with a Coq process
 32
 33Communication with a Coq process can be done either through a high-level or
 34low-level interface.
 35
 36## Highlevel interface
 37
 38The function `capnp_message_generator` converts a socket into a nested generator that
 39yields request messages for predictions and expects to be sent response messages
 40in return. A simple example of how to handle these messages is in `pytact.fake_python_server`.
 41Some docs can be found with `capnp_message_generator`.
 42
 43The `capnp_message_generator` function can also dump the sequence of messages send
 44and received to a file. One can then use `capnp_message_generator_from_file` to replay
 45that sequence against a server, either in a unit-test where the call to `capnp_message_generator`
 46is mocked by `capnp_message_generator_from_file`, or using a fully fledged socket test
 47as can be found in `pytact.fake_coq_client`.
 48
 49## Lowlevel interface
 50
 51The function `capnp_message_generator_lowlevel` converts a socket into a generator that
 52yields lowlevel Cap'n Proto request messages for predictions and expects to be sent Cap'n proto
 53messages in return. There are four types of messages `msg` Coq sends.
 541. Synchronize: If `msg.is_synchronize` is true, Coq is attempting to synchronize
 55   its state with the server. A `PredictionProtocol.Response.synchronized` message is expected
 56   in return.
 572. Initialize: If `msg.is_initialize` is true, Coq is sending a list of available
 58   tactics and a global context fragment to be added to an existing stack of global context
 59   information. An empty stack can be created through `empty_online_definitions_initialize`.
 60   To add an initialize message to the stack, you can use
 61   ```
 62   with pytact.data_reader.online_definitions_initialize(existing_stack, msg) as definitions:
 63       print(type(definitions))
 64   ```
 65   Any subsequent messages will be made in the context of the
 66   tactics and predictions sent in this message, until an initialize message is received
 67   such that `msg.initialize.stack_size` is smaller than the current stack size.
 68   A `PredictionProtocol.Response.initialized` message is expected in response of this message.
 694. Predict: If `msg.is_predict` is true, Coq is asking to predict a list of plausible
 70   tactics given a proof state. The proof state can be easily accessed using
 71   ```
 72   with pytact.data_reader.online_data_predict(definitions, msg.predict) as proof_state:
 73       print(dir(proof_state))
 74   ```
 75   A `PredictionProtocol.Response.prediction` or `PredictionProtocol.Response.textPrediction`
 76   message is expected in return.
 775. Check Alignment: If `msg.is_check_alignment` is true, then Coq is asking the server
 78   to check which tactics and definitions are known to it. A
 79   `PredictionProtocol.Response.alignment` message is expected in return.
 80
 81You can wrap the generator produced by `capnp_message_generator_lowlevel` into
 82`record_lowlevel_generator` in order to record the exchanged messages to disk.
 83This trace can later be replayed using
 84`capnp_message_generator_from_file_lowlevel` (or using the high-level interface).
 85
 86"""
 87
 88from __future__ import annotations
 89from contextlib import contextmanager, ExitStack
 90from dataclasses import dataclass
 91from typing import Any, Callable, TypeAlias, TypeVar, Union, cast, BinaryIO
 92from collections.abc import Iterable, Sequence, Generator
 93from pathlib import Path
 94from immutables import Map
 95import signal
 96import pytact.graph_api_capnp_cython as apic # type: ignore
 97import capnp# type: ignore
 98import pytact.graph_api_capnp as graph_api_capnp # type: ignore
 99import socket
100import mmap
101import itertools
102import resource
103import threading
104import subprocess
105import shutil
106import time
107
108@contextmanager
109def file_dataset_reader(fname: Path) -> Generator[apic.Dataset_Reader, None, None]:
110    """Load a single dataset file into memory, and expose its raw Cap'n Proto structures.
111
112    This is a low-level function. Prefer to use `data_reader` or `lowlevel_data_reader`.
113    """
114    ...
115
116class LowlevelDataReader:
117    """A thin wrapper around the raw Cap'n Proto structures contained in a dataset directory.
118
119    Every file in the directory is assigned an integer called a graph-id. This class is a
120    `Sequence` that allows the retrieval of the structures in a file by its graph-id.
121    Additionally, the function `local_to_global` translates a dependency-index relative to
122    a graph-id to a new graph-id. This allows one to find out in which file a node in a
123    graph is located.
124
125    This is a lowlevel interface. For details on using the exposed structures the documentation
126    in the Cap'n Proto API file.
127    """
128
129    graphid_by_filename : dict[Path, int]
130    """Map from filenames in the data directory to their graph-id's."""
131
132    graph_files : list[Path]
133    """Map's a graph-id to a filename. Inverse of `graphid_by_filename`."""
134
135    def local_to_global(self, graph: int, dep_index: int) -> int:
136        """Convert a dependency-index relative to a graph-id to a new graph-id.
137
138        This is used to find the graph-id where a particular node can be found. If `graph` is the
139        graph-id that contains a reference to a node and `dep_index` is the relative location
140        where that node can be found then `local_to_global(graph, dep_index)` finds the graph-id
141        of the file where the node is physically located.
142        """
143        ...
144
145    def __getitem__(self, graph: int):
146        """Retrieve the raw Cap'n Proto structures associated with a graph-id."""
147        ...
148
149    def __len__(self) -> int: ...
150
151@contextmanager
152def lowlevel_data_reader(dataset_path: Path) -> Generator[LowlevelDataReader, None, None]:
153    """Map a directory of dataset files into memory, and expose their raw Cap'n Proto structures.
154
155    Arguments:
156    `dataset_path`: Can either be a dataset directory, or a SquashFS image file. In case of an image
157                    file `dataset.squ`, the image will be mounted using `squashfuse` on directory
158                    `dataset/`, after which the contents of that directory will be read.
159
160    This is a low-level function. Prefer to use `data_reader`. See `LowlevelDataReader` for
161    further documentation.
162    """
163    ...
164
165ProofStateId = int
166TacticId = int
167GraphId = int
168NodeId = int
169NodeHash = int
170
171class Node:
172    """A node in the calculus of inductive construction graph.
173
174    A node has a `label`, a unique `identity` and `children`. Some nodes are also
175    a `Definition`.
176    """
177
178    graph : GraphId
179    """The id of the graph in which this node occurs. For lowlevel use only."""
180
181    nodeid : NodeId
182    """The local id of the node in the dataset. For lowlevel use only."""
183
184    def __repr__(self):
185        ...
186
187    def __eq__(self, other: object) -> bool:
188        """Physical equality of two nodes. Note that two nodes with a different physical equality
189        may have the same `identity`. See `identity` for details.
190        """
191        ...
192
193    def __hash__(self):
194        """A hash that is corresponds to physical equality, not to be confused by `identity`."""
195        ...
196
197    @property
198    def label(self) -> apic.Graph_Node_Label_Reader:
199        """The label of the node, indicating it's function in the CIC graph."""
200        ...
201
202    @property
203    def identity(self) -> NodeHash:
204        """The identity of a node uniquely determines it. That is, one can consider any to nodes with the same
205        identity to be equal. The notion of equal we use is as follows:
206        1. Graph perspective: Two nodes have the same identity if and only if they are bisimilar.
207           In this notion, bisimilarity does take into consideration the label of the nodes, modulo some
208           equivalence class that is not fully specified here. One aspect of the equivalence class is that
209           for definition nodes their associated global context (accessed through `Definition.previous`) is
210           not taken into account.
211        2. Lambda calculus perspective: Two nodes have the same identity if their corresponding lambda terms
212           are alpha-equivalent. Note that two definitions with the same body are not considered
213           alpha-equivalent.
214
215        The identity of a node is used to perform partial graph-sharing. That is, two nodes with the same
216        identity are merged when the graph is generated. There are two reasons why two nodes with the same
217        semantic identity might have a different physical identity:
218        1. Nodes are only merged when the first node exists in the same graph as the second node, or exists
219           in a dependent graph. Hence, nodes originating from developments that do not depend on each other
220           are never merged. Full graph-sharing would require global analysis on a dataset, which any consumer
221           can optionally do as a post-processing step.
222        2. Two definition nodes with the same name and body have the same identity. But if they occur in
223           different global contexts, these nodes are physically different to ensure the integrity of their
224           global contexts.
225
226        Beware that the identity is currently a 64 bit field. For datasets that have graphs of size in the
227        order of billions of nodes there is a non-trivial chance of a collision. (We consider this acceptable
228        for machine learning purposes.)
229        """
230        ...
231
232    @property
233    def children(self) -> Sequence[tuple[int, Node]]:
234        """The children of a node, together with the labels of the edges towards the children.
235        Note that for efficiency purposes, the label is represented as an integer. The corresponding
236        enum of this integer is `graph_api_capnp.EdgeClassification`."""
237        ...
238
239    @property
240    def definition(self) -> Definition | None:
241        """Some nodes in the CIC graph represent definitions. Such nodes contain extra information about the
242        definition.
243        """
244        ...
245
246class ProofState:
247    """A proof state represents a particular point in the tactical proof of a constant."""
248
249    @property
250    def lowlevel(self):
251        ...
252    def __repr__(self):
253        ...
254
255    @property
256    def root(self) -> Node:
257        """The entry-point of the proof state, all nodes that are 'part of' the proof state are
258        reachable from here."""
259        ...
260
261    @property
262    def context(self) -> Sequence[Node]:
263        """The local context of the proof state, given as a tuple of their original name as they appeared in the
264        proof and the corresponding node.
265
266        The nodes always have either label `contextAssum` or `contextDef`.
267        Note that these nodes are also reachable from the root of the proof state.
268
269        The names of the context elements should be used for debugging and viewing purposes only, because
270        hypothesis-generating tactics have been modified to use auto-generated names. Hence, tactics should
271        not be concerned about the names of the context.
272        """
273        ...
274
275    @property
276    def context_names(self) -> Sequence[str]:
277        """The names of the local context nodes of the proof state, as they originally appeared in the proof.
278
279        These names should be used for debugging and viewing purposes only, because hypothesis-generating tactics have
280        been modified to use auto-generated names. Hence, tactics should not be concerned about the names of
281        the context.
282        """
283        ...
284
285    @property
286    def context_text(self) -> Sequence[str]:
287        """A textual representation of the type/definition of context nodes
288        """
289        ...
290
291    @property
292    def conclusion_text(self) -> str:
293        """A textual representation of the conclusion of the proof state.
294        """
295        ...
296
297    @property
298    def text(self) -> str:
299        """A textual representation of the proof state."""
300        ...
301
302    def __str__(self) -> str:
303        ...
304
305    @property
306    def id(self) -> ProofStateId:
307        """
308        A unique identifier of the proof state. Any two proof states in a tactical proof that have an equal id
309        can morally be regarded to be 'the same' proof state.
310        IMPORTANT: Two proof states with the same id may still have different contents. This is because proof states
311                   can contain existential variables (represented by the `evar` node) that can be filled as a
312                   side-effect by a tactic running on another proof state.
313        """
314        ...
315
316Unresolvable : TypeAlias = None
317Unknown : TypeAlias = None
318class Outcome:
319    """An outcome is the result of running a tactic on a proof state. A tactic may run on multiple proof states."""
320
321    tactic : apic.Tactic_Reader | Unknown
322    """The tactic that generated the outcome. For it's arguments, see `tactic_arguments`
323
324    Sometimes a tactic cannot or should not be recorded. In those cases, it is marked as 'unknown'.
325    This currently happens with tactics that are run as a result of the `Proof with tac` construct and it
326    happens for tactics that are known to be unsafe like `change_no_check`, `fix`, `cofix` and more.
327    """
328
329    @property
330    def lowlevel(self):
331        ...
332    def __repr__(self):
333        ...
334
335    @property
336    def before(self) -> ProofState:
337        """The proof state before the tactic execution."""
338        ...
339
340    @property
341    def after(self) -> Sequence[ProofState]:
342        """The new proof states that were generated by the tactic."""
343        ...
344
345    @property
346    def term(self) -> Node:
347        """The proof term that witnesses the transition from the before state to the after states. It contains a
348        hole (an `evar` node) for each of the after states. It may also refer to elements of the local context of
349        the before state.
350        """
351        ...
352
353    @property
354    def term_text(self) -> str:
355        """A textual representation of the proof term."""
356        ...
357
358    @property
359    def tactic_arguments(self) -> Sequence[Node | Unresolvable]:
360        """The arguments of the tactic that produced this outcome.
361
362        The node is the root of a graph representing an argument that is a term in the calculus of constructions.
363        Sometimes, an argument is not resolvable to a term, in which case it is marked as `Unresolvable`.
364        """
365        ...
366
367class ProofStep:
368    """A proof step is the execution of a single tactic on one or more proof states, producing a list of outcomes.
369    """
370
371    @property
372    def lowlevel(self):
373        ...
374    def __repr__(self):
375        ...
376
377    @property
378    def tactic(self) -> apic.Tactic_Reader | Unknown:
379        """The tactic that generated the proof step. Note that the arguments of the tactic can be found in the
380        individual outcomes, because they may be different for each outcome.
381
382        Sometimes a tactic cannot or should not be recorded. In those cases, it is marked as 'unknown'.
383        This currently happens with tactics that are run as a result of the `Proof with tac` construct and it
384        happens for tactics that are known to be unsafe like `change_no_check`, `fix`, `cofix` and more.
385        """
386        ...
387
388    @property
389    def outcomes(self) -> Sequence[Outcome]:
390        """A list of transformations of proof states to other proof states, as executed by the tactic of the proof
391        step
392        """
393        ...
394
395@dataclass
396class Original:
397    """Token that signals that a definition is original, as inputted by the
398    user."""
399    pass
400
401@dataclass
402class Discharged:
403    """Token that signals that a definition is derived from an original after
404    section closing."""
405
406    original: Definition
407
408@dataclass
409class Substituted:
410    """Token that signals that a definition is derived from an original after
411    module functor instantiation."""
412
413    original: Definition
414
415@dataclass
416class Inductive:
417    """Token that signals that a definition is inductive."""
418
419    representative: Definition
420    """The representative of the mutually recursive cluster. For lowlevel use only"""
421
422@dataclass
423class Constructor:
424    """Token that signals that a definition is a constructor."""
425
426    representative: Definition
427    """The representative of the mutually recursive cluster. For lowlevel use only"""
428
429@dataclass
430class Projection:
431    """Token that signals that a definition is a projection."""
432
433    representative: Definition
434    """The representative of the mutually recursive cluster. For lowlevel use only"""
435
436@dataclass
437class ManualConstant:
438    """Token that signals that a definition was inputted by the user by
439    manually inputting a term."""
440    pass
441
442@dataclass
443class TacticalConstant:
444    """Token that signals that a lemma was proved by the user using
445    a tactic proof."""
446
447    proof: Sequence[ProofStep]
448    """The proof of the lemma."""
449
450@dataclass
451class ManualSectionConstant:
452    """Token that signals that a constant is a section variable or let-construct."""
453    pass
454
455@dataclass
456class TacticalSectionConstant:
457    """Token that signals that a constant is a section let-construct with a tactic proof."""
458
459    proof: Sequence[ProofStep]
460    """The proof of the lemma."""
461
462class Definition:
463    """A definition of the CIC, which is either a constant, inductive, constructor, projection or section
464    variable. Constants and section variables can have tactical proofs associated to them.
465    """
466
467    node : Node
468    """The node that is associated with this definition. It holds that
469    `definition.node.definition == definition`.
470    """
471
472    @property
473    def lowlevel(self):
474        ...
475    def __repr__(self):
476        ...
477
478    def __eq__(self, other: object) -> bool:
479        """Physical equality of two nodes. Note that two nodes with a different physical equality
480        may have the same `identity`. See `identity` for details.
481        """
482        ...
483
484    def __hash__(self):
485        """A hash that is corresponds to physical equality, not to be confused by `identity`."""
486        ...
487
488    @property
489    def name(self) -> str:
490        """The fully-qualified name of the definition. The name should be unique in a particular global context,
491        but is not unique among different branches of the global in a dataset.
492        """
493        ...
494
495    @property
496    def previous(self) -> Definition | None:
497        """The previous definition within the global context of the current file.
498
499        Note that this is a lowlevel property. Prefer to use `global_context` or `clustered_global_context`.
500
501        The contract on this field is that any definition nodes reachable from the forward closure of the definition
502        must also be reachable through the chain of previous fields. An exception to this rule are mutually
503        recursive definitions. Those nodes are placed into the global context in an arbitrary ordering.
504        """
505        ...
506
507    @property
508    def external_previous(self) -> Sequence[Definition]:
509        """A list of definitions that are the representatives of files other than the current file that are
510        part of the global context. This list becomes populated when a `Require` statement occurred right before
511        the definition.
512
513        Note that this is a lowlevel property. Prefer to use `global_context` or `clustered_global_context`.
514        """
515        ...
516
517    def global_context(self, across_files : bool = True, inclusive = False) -> Iterable[Definition]:
518        """All of the definitions in the global context when this definition was created.
519
520        Note that this does not include this definition itself, except when the definition is a inductive,
521        constructor or projection. Because those are mutually recursive objects, they reference themselves
522        and are therefore part of their own global context.
523
524        The resulting iterable is topologically sorted. That is, for any definition in the stream, any
525        definition reachable from the forward closure of the definition also exists in the remainder of the
526        stream. An exception to this rule are mutually recursive definitions, because no topological sort
527        is possible there (see also `clustered_global_context`).
528
529        Arguments:
530        * `across_files`: if `False`, outputs only the definitions from the local file, default `True`.
531        * `inclusive`: if `True`, outputs also itself, default `False`.
532        Note: if it is a self-recursive definition, the `inclusive` argument is ignored, and considered as `True`
533        """
534        ...
535
536    def clustered_global_context(self, across_files : bool = True, inclusive : bool = False) -> Iterable[list[Definition]]:
537        """All of the definitions in the global context when this definition was created, clustered into
538        mutually recursive cliques.
539
540        The resulting iterable is topologically sorted. That is, for any definition in the stream, any
541        definition reachable from the forward closure of the definition also exists in the remainder of the
542        stream.
543
544        Arguments:
545        * `across_files`: if `False`, outputs only the definitions from the local file, default `True`.
546        * `inclusive`: if `True`, outputs also the cluster of itself, default `False`.
547        Note: if it is a self-recursive definition, the `inclusive` argument is ignored, and considered as `True`.
548        """
549        ...
550
551    @property
552    def status(self) -> Original | Discharged | Substituted:
553        """
554        A definition is either
555        (1) an object as originally inputted by the user.
556        (2) a definition that was originally defined in a section and has now had the section
557            variables discharged into it.
558        (3) a definition that was obtained by performing some sort of module functor substitution.
559        When a definition is not original, we cross-reference to the definition that it was derived from.
560        """
561        ...
562
563    @property
564    def kind(self) -> Union[Inductive, Constructor, Projection,
565                            ManualConstant, TacticalConstant,
566                            ManualSectionConstant, TacticalSectionConstant]:
567        """The kind of the definition.
568
569        It can be an constant, section constant, inductive, constructor or projection. In case of a constant
570        of a constant or section constant, an optional tactical proof can be associated. In case of an inductive,
571        constructor or projection, another definition that acts as the representative of the mutually recursive
572        cluster of definitions is associated.
573
574        The associated information is low-level. Prefer to use the properties `proof` and `cluster` to access them.
575        """
576        ...
577
578    @property
579    def proof(self) -> Sequence[ProofStep] | None:
580        """An optional tactical proof that was used to create this definition."""
581        ...
582
583    @property
584    def cluster_representative(self) -> Definition:
585        """A unique representative of the mutually recursive cluster of definitions this definition is part of.
586        If the definition is not mutually recursive, the representative is itself.
587
588        This is a low-level property. Prefer to use the `cluster` property.
589        """
590        ...
591
592    @property
593    def cluster(self) -> Iterable[Definition]:
594        """The cluster of mutually recursive definitions that this definition is part of.
595        If the definition is not mutually recursive, the cluster is a singleton."""
596        ...
597
598    @property
599    def type_text(self) -> str:
600        """A textual representation of the type of this definition."""
601        ...
602
603    @property
604    def term_text(self) -> str | None:
605        """A textual representation of the body of this definition.
606        For inductives, constructors, projections, section variables and axioms this is `None`.
607        """
608        ...
609
610    @property
611    def is_file_representative(self) -> bool:
612        """Returns true if this definition is the representative of the super-global context of it's file.
613        Se also `Dataset.representative`."""
614        ...
615
616class Dataset:
617    """The data corresponding to a single Coq source file. The data contains a representation of all definitions
618    that have existed at any point throughout the compilation of the source file.
619    """
620
621    graph : GraphId
622    filename : Path
623    """The physical file in which the data contained in this class can be found."""
624
625    @property
626    def lowlevel(self):
627        ...
628    def __repr__(self):
629        ...
630
631    @property
632    def dependencies(self) -> Sequence[Path]:
633        """A list of physical paths of data files that are direct dependencies of this file.
634
635        It is guaranteed that no cycles exist in the dependency relation between files induced by this field.
636        """
637        ...
638
639    @property
640    def representative(self) -> Definition | None:
641        """The entry point of the global context of definitions that are available when this file is 'Required' by
642        another file. The full global context can be obtained by following the `previous` node of definitions.
643        If the compilation unit does not contain any 'super'-global definitions this may be `None`.
644
645        This is a low-level property.
646        Prefer to use `definitions(spine_only = True)` and `clustered_definitions(spine_only=True)`.
647        """
648        ...
649
650    def definitions(self, across_files = False, spine_only = False) -> Iterable[Definition]:
651        """All of the definitions present in the file.
652        Note that some of these nodes may not be part of the 'super-global' context. Those are definitions inside
653        of sections or module functors.
654
655        Arguments:
656        * `across_files`: if `True`, outputs also definitions from dependent files, default `False`.
657        * `spine_only`: if True, outputs only the definitions on the main spine, default `False`.
658        Note: `across_files = True` is incompatible with `spine_only = False`.
659
660        When `spine_only = True`, the resulting iterable is topologically sorted. That is, for
661        any definition in the stream, any definition reachable from the forward closure of the definition
662        also exists in the remainder of the stream. An exception to this rule are mutually recursive definitions,
663        because no topological sort is possible there (see also `clustered_definitions`).
664        """
665        ...
666
667    def clustered_definitions(self, across_files = False, spine_only = False) -> Iterable[list[Definition]]:
668        """All of the definitions present in the file, clustered by mutually recursive definitions.
669        Note that some of these nodes may not be part of the 'super-global' context. Those are definitions inside
670        of sections or module functors.
671
672        Arguments:
673        * `across_files`: if `True`, outputs also definitions from dependent files, default `False`.
674        * `spine_only`: if `True`, outputs only the definitions on the main spine, default `False`.
675        Note: `across_files = True` is incompatible with `spine_only = False`.
676
677        When `spine_only = True`, the resulting iterable is topologically sorted. That is, for
678        any definition in the stream, any definition reachable from the forward closure of the definition
679        also exists in the remainder of the stream.
680        """
681        ...
682
683    @property
684    def module_name(self) -> str:
685        ...
686
687    def node_by_id(self, nodeid: NodeId) -> Node:
688        """Lookup a node inside of this file by it's local node-id. This is a low-level function."""
689        ...
690
691def lowlevel_to_highlevel(lreader : LowlevelDataReader) -> dict[Path, Dataset]:
692    """Convert a dataset initialized as a `LowLevelDataReader` into a high level interface."""
693    ...
694
695@contextmanager
696def data_reader(dataset_path: Path) -> Generator[dict[Path, Dataset], None, None]:
697    """Load a directory of dataset files into memory, and expose the data they contain.
698    The result is a dictionary that maps physical paths to `Dataset` instances that allow access to the data.
699
700    Arguments:
701    `dataset_path`: Can either be a dataset directory, or a SquashFS image file. In case of an image
702                    file `dataset.squ`, the image will be mounted using `squashfuse` on directory
703                    `dataset/`, after which the contents of that directory will be read.
704
705    """
706    ...
707
708class OnlineDefinitionsReader:
709    """This class represents a global context to Python by Coq, containing all
710    known definitions and available tactics."""
711
712    def local_to_global(self, graph: int, dep_index: int) -> int:
713        """Convert a dependency-index relative to a graph-id to a new graph-id.
714
715        This is used to find the graph-id where a particular node can be found. If `graph` is the
716        graph-id that contains a reference to a node and `dep_index` is the relative location
717        where that node can be found then `local_to_global(graph, dep_index)` finds the index in
718        the global context stack where the node is physically located.
719
720        Lowlevel function.
721        """
722        ...
723
724    @property
725    def representative(self) -> Definition | None:
726        """The last definition in the global context. All other definitions can be accessed by following
727        the `Definition.previous` chain starting from this definitions.
728
729        This is a low-level property.
730        Prefer to use `definitions` and `clustered_definitions`.
731        """
732        ...
733
734    def definitions(self, full : bool = True) -> Iterable[Definition]:
735        """The list of definitions that are currently in the global context.
736
737        The resulting iterable is topologically sorted. That is, for any definition in the stream, any
738        definition reachable from the forward closure of the definition also exists in the remainder of
739        the stream. An exception to this rule are mutually recursive definitions,
740        because no topological sort is possible there (see also `clustered_definitions`).
741        """
742        ...
743
744    def clustered_definitions(self, full : bool = True) -> Iterable[list[Definition]]:
745        """All of the definitions present in the global context, clustered by mutually recursive definitions.
746
747        The resulting iterable is topologically sorted. That is, for any definition in the stream, any
748        definition reachable from the forward closure of the definition also exists in the remainder of
749        the stream.
750        """
751        ...
752
753@contextmanager
754def online_definitions_initialize(stack: OnlineDefinitionsReader,
755                                  init: apic.PredictionProtocol_Request_Initialize_Reader) -> Generator[OnlineDefinitionsReader, None, None]:
756    """Given a new initialization message sent by Coq, construct a
757    `OnlineDefinitiosnReader` object. This can be used to inspect the
758    definitions currently available. Additionally, using `online_data_predict`
759    it can be combined with a subsequent prediction request received from Coq
760    to build a `ProofState`.
761    """
762    ...
763
764def empty_online_definitions_initialize() -> OnlineDefinitionsReader:
765    ...
766
767@contextmanager
768def online_data_predict(base: OnlineDefinitionsReader,
769                        predict: apic.PredictionProtocol_Request_Predict_Reader) -> Generator[ProofState, None, None]:
770    """Given a `OnlineDefinitionsReader` instance constructed through
771    `online_data_initialize`, and a prediction message sent by Coq, construct a
772    `ProofState` object that represents the current proof state in Coq.
773    """
774    ...
775
776@dataclass
777class TacticPredictionGraph:
778    ident : int
779    arguments : list[Node]
780    confidence : float
781
782@dataclass
783class TacticPredictionsGraph:
784    predictions : list[TacticPredictionsGraph]
785
786@dataclass
787class TacticPredictionText:
788    tactic_text : str
789    confidence : float
790
791@dataclass
792class TacticPredictionsText:
793    predictions : list[TacticPredictionsText]
794
795@dataclass
796class CheckAlignmentMessage:
797    pass
798
799@dataclass
800class CheckAlignmentResponse:
801    unknown_definitions : list[Definition]
802    unknown_tactics : list[int]
803
804@dataclass
805class GlobalContextMessage:
806    """A message containing a new global context sent to Python from Coq,
807    including all known definitions and available tactics."""
808
809    definitions : OnlineDefinitionsReader
810    tactics : Sequence[apic.AbstractTactic_Reader]
811
812    log_annotation : str
813    """An annotation representing the current position of Coq in a source
814    document. For logging and debugging purposes."""
815
816    prediction_requests : Generator[GlobalContextMessage | ProofState | CheckAlignmentMessage,
817                                    None | TacticPredictionsGraph | TacticPredictionsText | CheckAlignmentResponse,
818                                    None]
819    """A sub-generator that produces new requests from Coq that are based on or
820    extend the global context of the current message. Once the sub-generator
821    runs out, the parent generator continues."""
822
823def capnp_message_generator_lowlevel(socket: socket.socket) -> (
824        Generator[apic.PredictionProtocol_Request_Reader,
825                  capnp.lib.capnp._DynamicStructBuilder, None]):
826    """A generator that facilitates communication between a prediction server and a Coq process.
827
828    Given a `socket`, this function creates a generator that yields messages of type
829    `pytact.graph_api_capnp_cython.PredictionProtocol_Request_Reader` after which a
830    `capnp.lib.capnp._DynamicStructBuilder` message needs to be `send` back.
831    """
832    ...
833
834def capnp_message_generator_from_file_lowlevel(
835        message_file: BinaryIO,
836        check : Callable[[Any, Any, Any], None] | None = None) -> (
837        Generator[apic.PredictionProtocol_Request_Reader,
838                  capnp.lib.capnp._DynamicStructBuilder, None]):
839    """Replay and verify a pre-recorded communication sequence between Coq and a prediction server.
840
841    Lowlevel variant of `capnp_message_generator_from_file`.
842
843    Accepts a `message_file` containing a stream of Capt'n Proto messages as recorded using
844    `capnp_message_generator(s, record=message_file)`. The resulting generator acts just like the generator
845    created by `capnp_message_generator` except that it is not connected to the socket but just replays the
846    pre-recorded messages.
847
848    Arguments:
849    - `check` is an optional callable that can be used to compare the response of the server to the recorded
850      response. It accepts three arguments: The recorded message that was sent by Coq, the response of the server
851      and the recorded response.
852    """
853    ...
854
855def record_lowlevel_generator(
856        record_file: BinaryIO,
857        gen: Generator[apic.PredictionProtocol_Request_Reader,
858                       capnp.lib.capnp._DynamicStructBuilder, None]) -> (
859                           Generator[apic.PredictionProtocol_Request_Reader,
860                                     capnp.lib.capnp._DynamicStructBuilder, None]):
861    """Record a trace of the full interaction of a lowlevel generator to a file
862
863    Wrap a lowlevel generator (such as from `capnp_message_generator_lowlevel`) and dump all exchanged messages
864    to the given file. The file can later be replayed with `capnp_message_generator_from_file_lowlevel`.
865    """
866    ...
867
868def prediction_generator(
869        lgenerator: Generator[apic.PredictionProtocol_Request_Reader,
870                              capnp.lib.capnp._DynamicStructBuilder, None],
871        defs: OnlineDefinitionsReader):
872    """Given the current global context stack `defs`, convert a low-level
873    generator to a high-level `GlobalContextMessage`"""
874    ...
875
876def capnp_message_generator(socket: socket.socket, record: BinaryIO | None = None) -> GlobalContextMessage:
877    """A generator that facilitates communication between a prediction server and a Coq process.
878
879    Given a `socket`, this function creates a `GlobalContextMessage` `context`. This message contains an
880    initially empty list of available tactics and definitions in the global context. Through
881    `context.prediction_requests` one can access a generator that yields prediction requests and expects
882    predictions to be sent in response. The possible messages are as follows:
883    - `GlobalContextMessage`: An additional, nested, global context message that amends the current context
884       with additional tactics and definitions. The prediction requests of this nested context need to be
885       exhausted before continuing with messages from the current context.
886    - `CheckAlignmentMessage`: A request to check which of Coq's current tactics and definitions the
887      prediction server currently "knows about". The generator expects a `CheckAlignmentResponse` to be
888      sent in response.
889    - `ProofState`: A request to make tactic predictions for a given proof state. Either a
890      `TacticPredictionsGraph` or a `TacticPredictionsText` message is expected in return.
891
892    When `record` is passed a file descriptor, all received and sent messages will be dumped into that file
893    descriptor. These messages can then be replayed later using `capnp_message_generator_from_file`.
894    """
895    ...
896
897def capnp_message_generator_from_file(message_file: BinaryIO,
898                                      check : Callable[[Any, Any, Any], None] | None = None,
899                                      record: BinaryIO | None = None) -> GlobalContextMessage:
900    """Replay and verify a pre-recorded communication sequence between Coq and a prediction server.
901
902    Highlevel variant of `capnp_message_generator_from_file_lowlevel`.
903
904    Accepts a `message_file` containing a stream of Capt'n Proto messages as recorded using
905    `capnp_message_generator(s, record=message_file)`. The resulting generator acts just like the generator
906    created by `capnp_message_generator` except that it is not connected to the socket but just replays the
907    pre-recorded messages.
908
909    Arguments:
910    - `check` is an optional callable that can be used to compare the response of the server to the recorded
911      response. It accepts three arguments: The recorded message that was sent by Coq, the response of the server
912      and the recorded response.
913    - `record` is an optional file that re-records the interaction with the current server
914    """
915    ...
916
917class GlobalContextSets:
918    """Lazily retrieve a the global context of a definition as a set, with memoization.
919
920    Because this class can allocate (large) amounts of memory, use it as a context-manager. To create a new
921    instance, and retrieve the context of a definition, use:
922    ```
923    with GlobalContextSets.new_context() as gcs:
924        gcs.global_context_set(my_definition)
925    ```
926    This context will be cached, as well as intermediate results (global contexts of other definitions) as
927    long as the context-set is in scope. Caching can be nested by using `sub_context`:
928    ```
929    with GlobalContextSets.new_context() as gcs:
930        gcs.global_context_set(my_definition)
931        with gcs.sub_context(lambda d:is_file_representative) as gcs_sub:
932            # gcs_sub remembers anything in the case of gcs
933            # caching is propagated to the parent cache only if the provided lambda evaluates to true
934            gcs_sub.global_context_set(my_definition2)
935    ```
936    """
937    def __init__(self, cache, parent: GlobalContextSets | None,
938                 should_propagate: Callable[[Definition], bool]):
939        """Do not call directly. Use `GlobalContextSets.new_context` and `GlobalContextSets.sub_context`."""
940        ...
941
942    def global_context_set(self, d: Definition) -> Map:
943        """Retrieve the global context of a definition, caching the result, and intermediate results."""
944        ...
945
946    @contextmanager
947    def sub_context(self, propagate: Callable[[Definition], bool]) -> Generator[GlobalContextSets, None, None]:
948        """Create a sub-context with a separate cache from it's parent, but sharing any info that was
949        already present in the parent's cache.
950
951        For any intermediate results for a definition `d` for which `propagate(d)` evaluates to true,
952        the result is propagated to the parent's cache."""
953        ...
954
955    @staticmethod
956    def new_context() -> Generator[GlobalContextSets, None, None]:
957        """Crate a new caching context where global-context-set's can be retrieved and cached."""
958        ...
959
960def node_dependencies(n: Node, deps: set[Definition] | None = None) -> set[Definition]:
961    """Given a `Node` `n`, calculate the set of `Definition`'s that are directly reachable
962    from `n`. Transitively reachable definitions are not included.
963    """
964    ...
965
966def definition_dependencies(d: Definition) -> set[Definition]:
967    """Given a `Definition` `d`, calculate the set of `Definition`'s that are directly reachable for `d`.
968    Transitively reachable definitions are not included nor is `d` itself.
969    """
970    ...