/*
    This magma script implements a special case of Algorithm 1 in the paper
    
       "Structure computation and discrete logarithms in finite abelian p-groups", A.V. Sutherland, Math. Comp. 80 (2011), 477-500
       
       http://www.ams.org/journals/mcom/2011-80-273/S0025-5718-10-02356-2/
       http://arxiv.org/abs/0809.3413
       
    for computing discrete logarithms in cyclic p-groups when p is small (e.g. 2 or 3).
    Anyone is welcome to use and/or modify this code however they see fit (but please cite the paper).
*/

// data structure for generic black box group, used to count operations and hide group implementation
BBGrp := recformat<group, p : Integers(), n : Integers(), gen, id, op, mults : Integers(), sqrs : Integers(), ppows : Integers()>;

// creates a black box cyclic group of prime power order generated by g (parent group may be larger)
function cyclic_pgroup (g)
    assert Type(g) eq GrpAbElt;
    isppower,p,n := IsPrimePower(Order(g));
    assert isppower;
    G := Parent(g);
    return rec<BBGrp | group := G, p := p, n := n, gen := g, id := Identity(G), op := func<g,h|G!g + G!h>, mults := 0, sqrs := 0, ppows := 0>;
end function;

// generic group operation, do not count operations involving the identity
function compose (G,g,h)
    if g eq G`id then return G,h; end if;
    if h eq G`id then return G,g; end if;
    if g eq h then G`sqrs +:= 1; else G`mults +:=1; end if;
    return G,G`op(g,h);
end function;

// given seq S of elements in generic group G returns their composition
function reduce (G,S)
    if #S eq 0 then return G,G`id; end if;
    a := S[1];
    for i:= 2 to #S do G,a:=compose(G,a,S[i]); end for;
    return G,a;
end function;

// returns a^p; assumes p is small (e.g. 2 or 3) (for larger p we would use an addition chain)
function ppower (G,a)
    // if G is a group in which p-powering can be optimized, this is where to do it
    if a eq G`id then return G,G`id; end if;
    b := G`group!a;
    for i:= 2 to G`p do _,b:=compose(G,b,a); end for;
    if G`p eq 2 then G`sqrs +:= 1; else G`ppows +:=1; end if;
    return G,b;
end function;

// returns [a,a^p,a^(p^2),a^(p^3),...,a^(p^(n-1))]
function ppowers (G,a,n)
    a := [G`group!a];
    for i in [1..n-1] do G,b:=ppower(G,a[#a]); a:=Append(a,b); end for;
    return G,a;
end function;

// simple implementation of Yao-style exponentiation using precomputed p-powers of base a
// note that this is only suitable for small p (e.g. p = 2 or 3)
// a should be an array with a[i+1] = a^(2^i) for i in [0..n-1]
// e is an integer exponent
function yao_exp (G,a,e)
    p:=G`p; n:=G`n;
    e:=e mod p^n;
    if e eq 0 then return G,G`id; end if;
    if e eq 1 then return G,a[1]; end if;
    // pick window size, assume we will never want to go past log_p(log_p(#G))
    bests:=0; bestc:=Infinity(); m:=0;
    for s:= 1 to Ceiling(Log(p,n)) do
        digits := IntegerToSequence(e,p^s);  m:=Max(digits);
        c := m-1+#[d:d in digits|d gt 0];
        if c lt bestc then bestc := c; bests := s; end if;
    end for;
    s:=bests; c:=bestc; digits:=IntegerToSequence(e,p^s); m:=Max(digits);
    A:=[[]:i in [1..m]];
    for i in [1..#digits] do d:=digits[i]; if d ne 0 then A[d]:=Append(A[d],i-1); end if; end for;
    // run the algorithm -- this is where all the group operations happen
    b := G`id; c := G`id;
    for d := m to 1 by -1 do
        G,c := reduce(G,[c]cat[a[s*i+1]:i in A[d]]);
        G,b := compose(G,b,c);
    end for;
    return G,b;
end function;

// given ppowers [a,a^p,a^(p^2),...,a^(p^n-1)] computes lookup table (associative array) for <a^(-p^(n-t))>
function table (G,a,t)
    p := G`p; a:=[a[i]:i in [#a-t+1..#a]];
    n := p^t;  assert n gt 1;
    A:=[G`id:i in [0..n-1]];
    for i:=0 to #a-1 do A[p^i] := a[i+1]; end for; // copy in pth powers we know
    for i:=2 to n-1 do
        if A[i] ne G`id then continue; end if; // don't recompute pth powers
        G,A[i] := compose(G,A[Floor(i/2)],A[Ceiling(i/2)]);
    end for;
    // sanity check, don't count this group operation
    _,b:=compose(G,A[n-1],A[1]); assert b eq G`id and A[G`p^(t-1)] ne G`id;
    T:=AssociativeArray(A);
    for i:=1 to #A do T[A[i]]:=(n-i) mod n; end for;
    return G,T;
end function;

// returns log_p q(j,k) as defined in paper
function qjk (j,k,n) return j+Max(0,n-k); end function;

// given p-powers [a,a^p,a^(p^2),...] returns a(j,k),a(j,k)^p,a(j,k)^(p^2,... as defined in paper
function ajkpowers (G,a,j,k,n) return [a[i+1]: i in [qjk(j,k,n)..#a-1]]; end function;

// DL base case (cyclic p-group)
// j,k identifies subgroup G(j,k):={g^(p^j):g^(p^k)=1:g in G}
//    since G has rank 1, this is just the unique cyclic subgroup of order p^(k-j)
// t >= k-j identifies base subgroup G(t,n) of order p^t
// T is a lookup table (associative array) for G(0,t)
function DLbase(G,j,k,b,t,T)
    x := T[b];
    // when G(j,k) is a proper subgroup of G(0,t) we need to remove powers of p from the result
    for i:=1 to t-(k-j) do x := ExactQuotient(x,G`p); end for;
    return G, x;
end function;
    
// Implmentation of Algorithm 1 in [1] specialized to cyclic p-groups
// a is an array of p-powers: a[i+1] = a^(p^i) for i in [0..n-1] where a is the *inverse* of our base
//   (alternatively, caller can negate the final result)
// n = log_p(|a|)
// j,k identifies subgroup G(j,k):={g^(p^j):g^(p^k)=1:g in G}
//    since G has rank 1, this is just the unique cyclic subgroup of order p^(k-j)
// b = element of <a> whose discrete log we want to compute
// t >= k-j identifies base subgroup G(0,t) of order p^t
// T is a lookup table (associative array) for DL's in G(0,t) (negated, so DLs wrt to 1/a)
// W is an optional list of hints for w (default w works reasonably well though)
function DL(G,a,n,j,k,b,t,T:W:=[])
    assert j le k and k le n and t gt 0;
    if k-j le t then return DLbase(G,j,k,b,t,T); end if;
    if #W gt 0 then
        w := W[1]; W:=Remove(W,1);
    else
        w := Ceiling(1.2*Log(2,k-j+1)/Log(2,Log(2,k-j+2)));
    end if;
    // make sure w is at least 2 and try to take advantage of precomputed base case table
    if w lt 2 then w:=2; else while w gt 2 and Ceiling((k-j)/(w-1)) le t do w-:=1; end while; end if;
    m := Floor((k-j)/w); d := (k-j)-m*w;
    J:=[j+i*m:i in [0..w]];
    for i:=1 to d do J[#J-d+i] +:= i; end for;
    G, g := ppowers (G, b, J[w]-j+1);
    g := [g[J[i]-j+1]:i in [1..w]];
    x:=0;
    for i := w to 1 by -1 do
        aa := ajkpowers(G,a,J[i],k,n);
        G,bb := yao_exp(G,aa,x);
        G,bb := compose(G,g[i],bb);
        G,v := $$(G,a,n,J[i],J[i+1],bb,t,T:W:=W);
        x +:= G`p^(qjk(J[i],J[i+1],n) - qjk(J[i],k,n))*v;
    end for;
    return G,x;
end function;
    
procedure TestDL(G,t,cnt:W:=[])
    G`mults:=0; G`sqrs:=0; G`ppows:=0;
    p := G`p; n := G`n;
    for i:=1 to cnt do
        G,a:=ppowers(G,G`gen,n);
        G,T:=table(G,a,t);
        e:=Random(p^n);
        _,b:=yao_exp(G,a,e);
        G,f:=DL(G,a,n,0,n,b,t,T:W:=W);
        assert IsDivisibleBy(e+f,p^n);
    end for;
    if p eq 2 then
        printf "Group order %o^%o, t=%o, used %o mults, %o squares per DL on average over %o tests\n", G`p, G`n, t, Round(G`mults/cnt), Round(G`sqrs/cnt), cnt;
    else
        printf "Group order %o^%o, t=%o, used %o mults, %o squares, %o %o-powers per DL on average over %o tests\n", G`p, G`n, t, Round(G`mults/cnt), Round(G`sqrs/cnt), Round(G`ppows/cnt), p, cnt;
    end if;
end procedure;

print "Computing 100 random discrete logs in mu(2^100)...";
TestDL(cyclic_pgroup(AdditiveGroup(Integers(2^100)).1),5,100:W:=[7,3]);
print "Computing 100 random discrete logs in mu(2^372)...";
TestDL(cyclic_pgroup(AdditiveGroup(Integers(2^372)).1),6,100);
print "Computing 100 random discrete logs in mu(3^239)...";
TestDL(cyclic_pgroup(AdditiveGroup(Integers(3^239)).1),4,100);
