Generics.mag

// This file is part of ExactpAdics
// Copyright (C) 2018 Christopher Doris
//
// ExactpAdics is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// ExactpAdics is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with ExactpAdics.  If not, see <http://www.gnu.org/licenses/>.


///# Internals
/// In this section, we document the generic internal structure common to all p-adic objects.
///toc
///
///## Overview
/// 
/// All p-adic objects are a subtype of `AnyPadExact`, which has attributes including:
/// - `id`: A global id, a unique integer assigned in dependency order.
/// - `approximation`: A list of approximations of the object, the `n`th entry being the approximation at the `n`th "epoch".
/// - `dependencies`: A list of the other p-adic object on which this depends.
/// - `get_approximation`: A function taking a positive integer `n` (an epoch) and the list of approximations of the dependencies at this epoch, and returning an approximation at this epoch.
/// - `min_epoch`: The lowest epoch for which `get_approximation` should be called. For lower epochs, interpolation is used (see `InterpolateEpochs`).
/// - `max_epoch`: The highest epoch for which `get_approximation` should be called. For higher epochs, a precision error is raised.
/// - `internal_data`: Some state which may be freely modified by `get_approximation`.
///
/// There are two direct subtypes of `AnyPadExact`: `StrPadExact` representing a structure (or set), and `PadExactElt` representing an element of such a structure. For example, `FldPadExact` represents a p-adic field and is a subtype of `StrPadExact`, and its elements have type `FldPadExactElt` which is a subtype of `PadExactElt`. `PadExactElt` has an additional `parent` attribute, which must be the parent structure to which the element belongs.

import "Utils.mag": Z, NEXTID, WEQ;

DEFAULT_MAX_EPOCH := 28;

///## Generic intrinsics

// represents any exact p-adic object (structure or element)
declare type AnyPadExact;
declare attributes AnyPadExact
  : id                   // global id
  , approximations       // list of approximations corresponding to each epoch
  , dependencies         // list of other AnyPadExact's this depends on
  , get_approximation    // function([*dependency approximations*]) -> approximation (or any callable)
  , min_epoch            // the first epoch we should call get_approximation on
  , max_epoch            // the last epoch we should call get_approximation on (not currently enforced)
  , internal_data        // internal data which can be freely modified by get_approximation
  ;

// exact p-adic structure
declare type StrPadExact[PadExactElt]: AnyPadExact;

// exact p-adic element
declare type PadExactElt: AnyPadExact;
declare attributes PadExactElt
  : parent               // parent
  ;

// valuation of a p-adic object
declare type ValPadExact;

procedure bring_to_epoch(x, n)
  if #x`approximations lt n then
    n := Max(n, x`min_epoch);
    error if assigned x`max_epoch and n gt x`max_epoch, "precision error: max_epoch exceeded";
    for d in x`dependencies do
      bring_to_epoch(d, n);
    end for;
    xds := [* d`approximations[n] : d in x`dependencies *];
    xx := GetApproximation(x`get_approximation, n, xds);
    SetApproximation(x, n, xx);
    assert #x`approximations eq n;
  end if;  
end procedure;

function can_bring_to_epoch(x, n)
  if #x`approximations lt n then
    n := Max(n, x`min_epoch);
    if assigned x`max_epoch and n gt x`max_epoch then
      return false;
    end if;
    for d in x`dependencies do
      if not can_bring_to_epoch(d, n) then
        return false;
      end if;
    end for;
    xds := [* d`approximations[n] : d in x`dependencies *];
    xx := GetApproximation(x`get_approximation, n, xds);
    SetApproximation(x, n, xx);
    assert #x`approximations eq n;
  end if;  
  return true;
end function;

intrinsic BringToEpoch(x :: AnyPadExact, n :: RngIntElt)
  {Brings x to epoch n.}
  bring_to_epoch(x, n);
end intrinsic;

intrinsic BringToEpoch(xs :: [AnyPadExact], n :: RngIntElt)
  {Brings each `x` in `xs` to epoch `n`.}
  for x in xs do
    bring_to_epoch(x, n);
  end for;
end intrinsic;

intrinsic CanBringToEpoch(x :: AnyPadExact, n :: RngIntElt) -> BoolElt
  {Tries to bring x to epoch n. Returns `true` if successful, or `false` if the max_epoch of something is exceeded.}
  return can_bring_to_epoch(x, n);
end intrinsic;

intrinsic CanBringToEpoch(xs :: [AnyPadExact], n :: RngIntElt) -> BoolElt
  {Tries to bring each `x` in `xs` to epoch `n`. Returns `true` if successful, or `false` if the max_epoch of something is exceeded.}
  for x in xs do
    if not can_bring_to_epoch(x, n) then
      return false;
    end if;
  end for;
  return true;
end intrinsic;

intrinsic EpochApproximation(xs :: [AnyPadExact], n :: RngIntElt) -> []
  {Returns sequence of approximations of xs at epoch n.}
  BringToEpoch(xs, n);
  return [x`approximations[n] : x in xs];
end intrinsic;

intrinsic EpochApproximation(x :: AnyPadExact, n :: RngIntElt) -> .
  {The approximation of x at epoch n.}
  BringToEpoch(x, n);
  return x`approximations[n];
end intrinsic;

intrinsic Init_AnyPadExact(x :: AnyPadExact)
  {Initializes x.}
  assert assigned x`dependencies;
  assert assigned x`get_approximation;
  if not assigned x`min_epoch then
    x`min_epoch := 1;
  else
    assert x`min_epoch ge 1;
  end if;
  x`id := NEXTID();
  assert forall{d : d in x`dependencies | d`id lt x`id};
  if not assigned x`max_epoch and #x`dependencies eq 0 then
    x`max_epoch := DEFAULT_MAX_EPOCH;
  end if;
  x`approximations := [**];
end intrinsic;

intrinsic Init_PadExactElt(x :: PadExactElt)
  {"}
  assert assigned x`parent;
  Init_AnyPadExact(x);
end intrinsic;

intrinsic Init(x :: AnyPadExact)
  {"}
  Init_AnyPadExact(x);
end intrinsic;

intrinsic Init(x :: PadExactElt)
  {"}
  Init_PadExactElt(x);
end intrinsic;

intrinsic Parent(x :: PadExactElt) -> StrPadExact
  {The parent of x.}
  return x`parent;
end intrinsic;

intrinsic BestApproximation(x :: AnyPadExact) -> .
  {The best approximation to x.}
  if #x`approximations eq 0 then
    BringToEpoch(x, 1);
  end if;
  return x`approximations[#x`approximations];
end intrinsic;

///hide
intrinsic 'eq'(E :: StrPadExact, F :: StrPadExact) -> RngIntElt
  {Equality.}
  return IsIdentical(E, F);
end intrinsic;

intrinsic ExistsEpochWith(x :: AnyPadExact, pred) -> BoolElt, RngIntElt
  {True if there is an epoch of x such that `pred(x)` is true. If so, returns it.}
  if #x`approximations eq 0 then
    if not CanBringToEpoch(x, 1) then
      return false, _;
    end if;
  end if;
  while not pred(x) do
    if not CanBringToEpoch(x, #x`approximations+1) then
      return false, _;
    end if;
  end while;
  return true, #x`approximations;
end intrinsic;

intrinsic ExistsEpochWithApproximation(x :: AnyPadExact, pred : FromTop:=false, Minimize:=false) -> BoolElt, RngIntElt
  {True if there is an epoch of x such that `pred(xx)` is true, where `xx` is the corresponding approximation. If so, returns it. False if `max_epoch` is ever exceeded.}
  if #x`approximations eq 0 then
    if not CanBringToEpoch(x, 1) then
      return false, _;
    end if;
  end if;
  nmin := FromTop select #x`approximations else 1;
  n := nmin;
  while not pred(x`approximations[n]) do
    n +:= 1;
    if not CanBringToEpoch(x, n) then
      return false, _;
    end if;
  end while;
  if Minimize and n eq nmin then
    while n gt 1 and pred(x`approximations[n-1]) do
      n -:= 1;
    end while;
  end if;
  return true, n;
end intrinsic;

intrinsic FirstEpochWithApproximation(x :: AnyPadExact, pred : FromTop:=false) -> RngIntElt
  {The first epoch whose approximation satisfies the predicate.}
  ok, n := ExistsEpochWithApproximation(x, pred : FromTop:=FromTop, Minimize);
  error if not ok, "precision error";
  return n;
end intrinsic;

intrinsic IncreaseEpochUntil(x :: AnyPadExact, pred)
  {Increases the epoch of x until `pred(x)` is true.}
  ok := ExistsEpochWith(x, pred);
  error if not ok, "precision error";
end intrinsic;

intrinsic IncreaseAbsolutePrecision(x :: PadExactElt, apr)
  {Increases the absolute precision of x to at least apr.}
  ok, a := IsValidAbsolutePrecisionDiff(x, apr);
  require ok: "apr is not a valid absolute precision for x";
  IncreaseEpochUntil(x, func<x | AbsolutePrecision(x) ge a>);
end intrinsic;

intrinsic Approximation(x :: PadExactElt, apr) -> .
  {An approximation of x with absolute precision at least apr.}
  ok, a := IsValidAbsolutePrecisionDiff(x, apr);
  require ok: "apr is not a valid absolute precision for x";
  n := FirstEpochWithApproximation(x, func<xx | AbsolutePrecision(xx) ge a> : FromTop);
  return x`approximations[n];
end intrinsic;

intrinsic InterpolateUpTo(x :: AnyPadExact, n :: RngIntElt)
  {Replaces the approximations of x below n with their interpolations.}
  require n ge 0: "n must be at least 0";
  if n eq 0 then
    return;
  end if;
  require n le #x`approximations: "n must be at most the current epoch";
  apprs := InterpolateEpochs(x, 0, n, x`approximations[n]);
  assert #apprs eq n-1;
  for i in [1..#apprs] do
    SetApproximation(x, i, apprs[i]);
  end for;
end intrinsic;

intrinsic InterpolateUpToFirstEpochWithApproximation(x :: AnyPadExact, pred : FromTop:=false) -> RngIntElt, RngIntElt
  {Replaces the approximations of x below n with their interpolations, where n is the first epoch whose approximation satisfies the predicate. Returns the new and old values of n.}
  n := FirstEpochWithApproximation(x, pred : FromTop:=FromTop);
  InterpolateUpTo(x, n);
  return FirstEpochWithApproximation(x, pred : FromTop:=FromTop), n;
end intrinsic;

/// Sets the approximation of x at epoch n to xx.
/// 
/// This is called by [`BringToEpoch`](#BringToEpoch) and [`InterpolateUpTo`](#InterpolateUpTo).
/// 
/// It checks the new approximation is valid by calling [`IsValidApproximation`](#IsValidApproximation), then calls [`SetApproximationHook`](#SetApproximationHook), then calls [`InterpolateEpochs`](#InterpolateEpochs) if necessary to get missing approximations below `n`, and then finally updates the list of approximations.
intrinsic SetApproximation(x :: AnyPadExact, n :: RngIntElt, xx)
  {Sets the approximation of x at epoch n to xx.}
  ok, err := IsValidApproximation(x, n, xx);
  require ok: "invalid approximation: " cat err;
  SetApproximationHook(x, n, ~xx);
  if n gt #x`approximations then
    if n gt #x`approximations + 1 then
      apprs := InterpolateEpochs(x, #x`approximations, n, xx);
      assert #apprs eq n - #x`approximations - 1;
      for appr in apprs do
        SetApproximation(x, #x`approximations + 1, appr);
      end for;
    end if;
    assert #x`approximations eq n-1;
    Append(~x`approximations, xx);
    assert #x`approximations eq n;
  else
    x`approximations[n] := xx;
  end if;
end intrinsic;

// inserts into deps the dependency [*x, Max(n,x`min_epoch)*] and recurses
procedure _backtrack(~deps, x, n, terminal_ids)
  // insert the dependency, if required
  ok, dep := IsDefined(deps, x`id);
  if not ok then
    n := Max(n, x`min_epoch);
    deps[x`id] := [* x, n *];
  elif dep[2] lt n then
    deps[x`id][2] := n;
  else
    return;
  end if;
  // if we get to here, we added a new dependency, so recurse
  if x`id notin terminal_ids then
    for d in x`dependencies do
      _backtrack(~deps, d, n, terminal_ids);
    end for;
  end if;
end procedure;

// inserts into deps the dependency x and recurses
procedure _backtrack_simple(~deps, x, terminal_ids)
  // insert the dependency, if required
  if IsDefined(deps, x`id) then
    return;
  else
    deps[x`id] := x;
  end if;
  // if we get to here, we added a new dependency, so recurse
  if x`id notin terminal_ids then
    for d in x`dependencies do
      _backtrack_simple(~deps, d, terminal_ids);
    end for;
  end if;
end procedure;

function backtrack(x, terminal_ids : simple:=false)
  deps := AssociativeArray(Z);
  if simple then
    _backtrack_simple(~deps, x, terminal_ids);
  else
    _backtrack(~deps, x, 0, terminal_ids);
  end if;
  return [* deps[id] : id in Sort(Setseq(Keys(deps))) *];
end function;

declare verbose ExactpAdics_WithDependencies, 1;

/// A copy of x with the direct dependencies ds.
/// 
/// This is useful if x is some complex expression in ds and you don't care about the intermediate expressions. It can provide a small speed-up in computing approximations for x.
/// 
///param Fast:=false When true, this in effect eliminates the intermediate variables entirely, often resulting in a large speed-up in computing approximations for x. It should not be used when there is an intermediate variable which may alter its cached approximation, since this cache is not considered. This includes operations such as division and polynomial discriminant; many such operations have a `Safe` parameter which makes them safe to use in this context. Note that the `min_epoch` of the copy of x will be the maximum of the `min_epoch` of the intermediate variables; similarly the `max_epoch` will be the minimum over intermediate variables.
intrinsic WithDependencies(x :: AnyPadExact, ds :: List : Fast:=false) -> AnyPadExact
  {A copy of x with the direct dependencies ds.}
  require forall{d : d in ds | ISA(Type(d), AnyPadExact) and d`id le x`id}: "ds must be dependencies of x";
  if [Z| d`id : d in x`dependencies] eq [Z| d`id : d in ds] then
    vprint ExactpAdics_WithDependencies: "WithDependencies: same dependencies";
    // special case: same list of dependencies
    return x;
  end if;
  z := UninitializedCopy(x);
  z`dependencies := ds;
  i := Index([Z| d`id : d in ds], x`id);
  if i ne 0 then
    vprint ExactpAdics_WithDependencies: "WithDependencies: x in ds";
    // special case: x in ds
    z`get_approximation := func<n, xds | xds[i]>;
  elif Fast then
    vprint ExactpAdics_WithDependencies: "WithDependencies: Fast";
    // precompute dependencies back to ds
    deps := backtrack(x, {d`id : d in ds} : simple);
    vprint ExactpAdics_WithDependencies: "WithDependencies: deps =", deps;
    assert deps[#deps]`id eq x`id;
    // for each dependency not in the original list, make a step consisting of the indices of its dependencies and its get_approximation function
    dids := {Z| d`id : d in ds};
    assert x`id notin dids;
    idxs := AssociativeArray(Z);
    for i in [1..#ds] do
      idxs[ds[i]`id] := i;
    end for;
    steps := [**];
    for d in deps do
      if d`id notin dids then
        Append(~steps, [* [* idxs[d2`id] : d2 in d`dependencies *], d`get_approximation *]);
        idxs[d`id] := #ds + #steps;
      end if;
    end for;
    vprint ExactpAdics_WithDependencies: "WithDependencies: steps =", steps;
    // this executes the steps in turn, to extend the list of approximations up to x
    z`get_approximation := function (n, xds)
      for step in steps do
        idxs, ga := Explode(step);
        xds2 := [* xds[i] : i in idxs *];
        xd := GetApproximation(ga, n, xds2);
        Append(~xds, xd);
      end for;
      return xds[#xds];
    end function;
    // aggregate the min_epoch and max_epoch of all dependencies
    z`min_epoch := Max([Z| d`min_epoch : d in deps]);
    max_epochs := [Z| d`max_epoch : d in deps | assigned d`max_epoch];
    if #max_epochs ne 0 then
      z`max_epoch := Min(max_epochs);
    end if;
  elif {Z| d`id : d in x`dependencies} subset {Z| d`id : d in ds} then
    vprint ExactpAdics_WithDependencies: "WithDependencies: more dependencies";
    // special case: expanding the set of dependencies
    dids := [Z| d`id : d in ds];
    idxs := [Z| Index(dids, d`id) : d in x`dependencies];
    xga := x`get_approximation;
    z`get_approximation := func<n, xds | GetApproximation(xga, n, [* xds[i] : i in idxs *])>;
    z`min_epoch := x`min_epoch;
    if assigned x`max_epoch then
      z`max_epoch := x`max_epoch;
    end if;
  else
    vprint ExactpAdics_WithDependencies: "WithDependencies: default";
    // we precompute the dependencies back to ds, and use this in get_approximation to update them all in a single forward pass
    deps := backtrack(x, {d`id : d in ds});
    vprint ExactpAdics_WithDependencies: "WithDependencies: deps =", deps;
    z`get_approximation := function (n, xds)
      for dep in deps do
        d, n0 := Explode(dep);
        n2 := Max(n0, n);
        if #d`approximations lt n2 then
          error if assigned d`max_epoch and n2 gt d`max_epoch, "precision error: max_epoch exceeded";
          xds2 := [* d2`approximations[n2] : d2 in d`dependencies *];
          xd2 := GetApproximation(d`get_approximation, n2, xds2);
          SetApproximation(d, n2, xd2);
        end if;
      end for;
      return x`approximations[n];
    end function;
  end if;
  Init(z);
  return z;
end intrinsic;

///## Overloadable intrinsics
///
/// The following intrinsics are expected to be overloaded for type-specific behaviour.

/// Assuming we have an approximation for `x` at epoch `n1`, and `xx2` is the approximation at epoch `n2`, returns approximations for the intermediate epochs `[n1+1..n2-1]`. Usually this will do something simple such as just lowering the precision of `xx2` to match each intermediate epoch.
intrinsic InterpolateEpochs(x :: AnyPadExact, n1 :: RngIntElt, n2 :: RngIntElt, xx2) -> List
  {The approximations for epochs [n1+1..n2-1].}
  if n1+1 ge x`min_epoch then
    return [* GetApproximation(x`get_approximation, n, [*d`approximations[n] : d in x`dependencies*]) : n in [n1+1..n2-1] *];
  end if;
  error "not implemented: InterpolateEpochs:", Type(x);
end intrinsic;

/// Given a `get_approximation` value `get`, an epoch `n` and a list of approximations of dependencies at epoch `n`, returns an approximation at epoch `n`.
intrinsic GetApproximation(get, n :: RngIntElt, xds :: List) -> .
  {Gets an approximation using get.}
  error "not implemented: GetApproximation:", Type(get);
end intrinsic;

/// Currently, `get_approximation` is always just a function, and so this just calls it.
intrinsic GetApproximation(get :: UserProgram, n :: RngIntElt, xds :: List) -> .
  {"}
  return get(n, xds);
end intrinsic;

/// Overload this to perform some action when we are about to set the approximation of `x` at epoch `n` to `xx`.
/// 
/// For example, this is used by `RngUPol_FldPadExact` to ensure all its approximations have the same variable name.
intrinsic SetApproximationHook(x :: AnyPadExact, n :: RngIntElt, ~xx)
  {Called by `SetApproximation`.}
  return;
end intrinsic;

intrinsic IsValidApproximation(x :: AnyPadExact, n :: RngIntElt, xx) -> BoolElt, MonStgElt
  {True if xx is a valid approximation of x at epoch n. If not, also returns an error message.}
  return true, _;
end intrinsic;

/// For p-adic elements, this checks that the approximation is an element of the approximation of the parent, and that it is weakly equal to the current best approximation for `x`.
intrinsic IsValidApproximation(x :: PadExactElt, n :: RngIntElt, xx) -> BoolElt, MonStgElt
  {"}
  if Parent(xx) cmpne EpochApproximation(Parent(x), n) then
    return false, "wrong parent";
  elif #x`approximations ne 0 and not WEQ(BestApproximation(x), xx) then
    return false, "inconsistent";
  else
    return true, _;
  end if;
end intrinsic;

// An uninitialized copy of x. Default implementation raises a not implemented error.
intrinsic UninitializedCopy(x :: AnyPadExact) -> AnyPadExact
  {An uninitialized copy of x.}
  error "not implemented: UninitializedCopy:", Type(x);
end intrinsic;

/// For structures, we raise a "cannot copy p-adic structures" error, since equality of structures is defined by equality of id, and this would break it.
intrinsic UninitializedCopy(x :: StrPadExact) -> StrPadExact
  {"}
  // We disallow copying of p-adic structures, because equality is defined in terms of id, and the copy would have a different id
  error "cannot copy p-adic structures";
end intrinsic;

/// For elements, we just need to copy the parent.
intrinsic UninitializedCopy(x :: PadExactElt) -> PadExactElt
  {"}
  z := New(Type(x));
  z`parent := x`parent;
  return z;
end intrinsic;