(* ::Package:: *)

(* :Title: Discrete Logarithms *)

(* :Author: Gerald Teschl 
			http://www.mat.univie.ac.at/~gerald/
			Gerald.Teschl@univie.ac.at  *)

(* :Summary:
	This package provides Pohlig-Hellman reduction together with the baby-step/giant-step algorithm for solving discrete log problems. *)

(* :Context: LogMod` *)

(* :Package Version: 1.2 *)

(* :Copyright: Copyright 2023, G. Teschl  *)

(* :History: 
	Version 1.0-1.2 by G. Teschl *)

(* :Keywords: Discrete log problem (DLP), Pohlig-Hellman reduction, Pollard's rho method, Shanks' baby-step/giant-step method, index calculus *)

(* :Source: None. *)

(* :Warning: None. *)

(* :Mathematica Version: 6.0 *)

(* :Limitation: The algortihms are not sutiable for handling large numbers. *)

(* :Discussion: This package computes the discrete logarithm.
    *)


BeginPackage["LogMod`"];

LogMod::usage = 
	"LogMod[a,b,p] computes the discrete logarithm of b with base a modulo p.";
LogModRho::usage =
	"LogModRho[a,b,p] computes the discrete logarithm of b with base a modulo p using Pollard's rho method.";
LogModBSGS::usage =
	"LogModBSGS[a,b,p] computes the discrete logarithm of b with base a modulo p using Shank's baby-step/giant-step method.";
LogModIC::usage =
	"LogModBSGS[a,b,p] computes the discrete logarithm of b with base a modulo p using the index calculus method.";

    


Begin["`Private`"];

LogMod::nosol="Warning! `2` is not a power of `1` modulo `3`";
LogMod::noinv="Warning! `1` is not invertible modulo `2`";
LogMod::method="Error! Unkonw method.";
LogMod::failedrho="Failed! Could not find a working collission.";
LogMod::failedic="Failed! Could not find sufficiently many relations.";

Options[LogMod] = {Method -> "Automatic", Verbose -> 0};

(* Choose the method *)

LogModSwitch[a_Integer,b_Integer,p_Integer,m_Integer,opts:OptionsPattern[LogMod]]:=Switch[OptionValue[Method],
	"Automatic",If[m<10^2,LogModNaive[a,b,p,m,opts],LogModBSGS[a,b,p,m,opts]],
	"Naive",LogModNaive[a,b,p,m,opts],
	"BSGS",LogModBSGS[a,b,p,m,opts],
	"IC",LogModIC[a,b,p,m,opts],
	"Rho",If[m<5,LogModNaive[a,b,p,m,opts],LogModRho[a,b,p,m,opts]], (* Pollard rho tends to fail for small values *)
	_,Message[LogMod::method];Null];

(* Naive method for small numbers *)
       
LogModNaive[a_Integer,b_Integer,p_Integer,m_Integer,opts:OptionsPattern[LogMod]]:=Block[{j=0,bb=Mod[b,p]},
	If[OptionValue[Verbose]>1,Print["LogModNaive: a=",a," b=",b," m=",m]];
	While[PowerMod[a,j,p]!=bb&&j<m,j++];
	If[j<m,j,Message[LogMod::nosol,a,b,p];Null]
];
LogModNaive[a_Integer,b_Integer,p_Integer,opts:OptionsPattern[LogMod]]:=LogModNaive[a,b,p,p-1,opts];

(* Pollard rho algorithm *)

LogModRho[a_Integer,b_Integer,p_Integer,m_Integer,opts:OptionsPattern[LogMod]]:=Block[{trys=10,j0,j=1,maxj=Floor[10Sqrt[m]],f,x0,y0,x,\[Alpha]x,\[Beta]x,y,\[Alpha]y,\[Beta]y,d,mm, pow, lam},
	If[OptionValue[Verbose]>1,Print["LogModRho: a=",a," b=",b," m=",m]];
	f[{x_,\[Alpha]_,\[Beta]_}]:=Block[{r=Mod[x,3]},Switch[r,
		0,{Mod[x a,p],Mod[\[Alpha]+1,m],\[Beta]},
		1,{PowerMod[x,2,p],Mod[2\[Alpha],m],Mod[2\[Beta],m]},
		2,{Mod[x b,p],\[Alpha],Mod[\[Beta]+1,m]}]];
	While[trys>0,
		If[OptionValue[Verbose]>0,Print["LogModRho: Searching for a collission"]];
		(*  Floyd's cycle detection algorithm
		j0=RandomInteger[{1,m-1}];
		x0=Mod[b PowerMod[a,j0,p],p];y0=x0;
		{x,\[Alpha]x,\[Beta]x}=f[{x0,j0,1}];{y,\[Alpha]y,\[Beta]y}=f[{x,\[Alpha]x,\[Beta]x}];
		While[j<maxj && x!=y,
			{x,\[Alpha]x,\[Beta]x}=f[{x,\[Alpha]x,\[Beta]x}];
			{y,\[Alpha]y,\[Beta]y}=f[f[{y,\[Alpha]y,\[Beta]y}]];
			j++
		];
		*)
		(* Brent's cycle detection algorithm. *)
		pow = 1; \[Lambda] = 1;
		j0=RandomInteger[{1,m-1}];
		{x,\[Alpha]x,\[Beta]x}={Mod[b PowerMod[a,j0,p],p],j0,1};
		{y,\[Alpha]y,\[Beta]y}=f[{x,\[Alpha]x,\[Beta]x}];
		While[j<maxj && x!=y,
			If[pow==\[Lambda],{x,\[Alpha]x,\[Beta]x}={y,\[Alpha]y,\[Beta]y};pow=2 pow;\[Lambda]=0];
			{y,\[Alpha]y,\[Beta]y}=f[{y,\[Alpha]y,\[Beta]y}];
			\[Lambda]++;
			j++
		];
		If[OptionValue[Verbose]>0,If[x==y,Print["LogModRho: Collission after ",j," steps."],Print["LogModRho: No collission after ",j," steps."]]]
		If[x==y&&Mod[\[Beta]y-\[Beta]x,m]!=0,
			d=GCD[\[Beta]y-\[Beta]x,m];mm=m/d;
			If[Mod[\[Alpha]x-\[Alpha]y,d]!=0,Message[LogMod::nosol,a,b,p];Return[Null,Block]];
			l=Mod[(\[Alpha]x-\[Alpha]y)/d PowerMod[(\[Beta]y-\[Beta]x)/d,-1,mm],mm];
			While[PowerMod[a,l,p]!=b,l=l+mm];
			Return[l,Block]
			];
		trys--;
	];
	Message[LogMod::failedrho];
	Null
];
LogModRho[a_Integer,b_Integer,p_Integer,opts:OptionsPattern[LogMod]]:=LogModRho[a,b,p,MultiplicativeOrder[a,p],opts];

(* Shanks baby-step/giant-step algorithm *)

LogModBSGS[a_Integer,b_Integer,p_Integer,q_Integer,opts:OptionsPattern[LogMod]]:=Block[{m=Ceiling[Sqrt[q]],bsteps=CreateDataStructure["HashTable"],bs=a,c,gs=b},
	If[OptionValue[Verbose]>1,Print["LogModBSGS: a=",a," b=",b," m=",m]];
	If[OptionValue[Verbose]>0,Print["LogModBSGS: Computing baby steps."]];
	bsteps["Insert",1->0];
	bsteps["Insert",a->1];
	Do[bs=Mod[a bs,p];bsteps["Insert",bs->j],{j,2,m-1}];
	If[OptionValue[Verbose]>0,Print["LogModBSGS: Computing giant steps."]];
	c=PowerMod[a,-m,p];
	Do[
		If[bsteps["KeyExistsQ",gs],Return[bsteps["Lookup",gs]+m j,Block]];
		gs=Mod[gs c,p]
	,{j,0,m-1}]
	Message[LogMod::nosol,a,b,p];
];
LogModBSGS[a_Integer,b_Integer,p_Integer,opts:OptionsPattern[LogMod]]:=LogModBSGS[a,b,p,MultiplicativeOrder[a,p],opts];

(* Index calculus *)

PrimeList[S_Real]:=Block[{primes={2},p=2},
	While[(p=NextPrime[p])<=S,primes=Append[primes,p]];
	primes];

LogModIC[a_Integer,b_Integer,p_Integer,m_Integer,opts:OptionsPattern[LogMod]]:=Block[{trys,B,myFactor,Mat={},Matt,f,j},
	If[OptionValue[Verbose]>1,Print["LogModIC: a=",a," b=",b," m=",m]];
	B=Prepend[PrimeList[Max[3,Exp[1./2 Sqrt[Log[p]Log[Log[p]]]]]],-1];
	If[OptionValue[Verbose]>0,Print["LogModIC: Using ",Length[B]," primes."]];
	trys=3m+5;
	myFactor[x_]:=Block[{K=Table[0,Length[B]],xx=x},
		If[x<0,K[[1]]=1;x=-x];
		Do[
			While[Mod[xx,B[[i]]]==0,K[[i]]++;xx=xx/B[[i]]],{i,2,Length[B]}];
			If[xx==1,K,{}]
		];
	If[OptionValue[Verbose]>0,Print["LogModIC: Searching for relations."]];
	Do[
		j=RandomInteger[m-1];
		f=myFactor[Mod[b a^j,p]];
		If[f!={},
			f=Mod[Append[Prepend[f,-1],j],m];
			Matt=Append[Mat,f];
			Matt=RowReduce[Matt,Modulus->m];
			If[Head[Matt]==List,Mat=Matt]; (* Check if RowReduce failed *)
			If[Norm[Mat[[-1]]]==0,Mat=Mat[[1;;-2]]];
			If[OptionValue[Verbose]>0,Print["LogModIC: New relation found: ",Length[Mat]]];
			If[Norm[Mat[[1,2;;-2]]]==0,Return[Mat[[1,-1]],Block]]
		];
	,{trys}];
	Message[LogMod::failedic];
	Null
];
LogModIC[a_Integer,b_Integer,p_Integer]:=LogModIC[a,b,p,MultiplicativeOrder[a,p]];

(* Pohlig-Hellman reduction for prime powers q^k *)

LogModPH[a_Integer,b_Integer,p_Integer,q_Integer,k_Integer,opts:OptionsPattern[LogMod]]:=Block[{a1,aj,bj,xj,yj},
	If[OptionValue[Verbose]>0,Print["LogModPH: a=",a," b=",b," q=",q," k=",k]];
	If[k==1||q^k<10^4,Return[LogModSwitch[a,b,p,q^k,opts]]];
	aj=PowerMod[a,q^(k-1),p];a1=aj;
	bj=PowerMod[b,q^(k-1),p];
	xj=LogModSwitch[a1,bj,p,q,opts];
	Do[
		aj=PowerMod[a,q^(k-j),p];
		bj=PowerMod[b,q^(k-j),p];
		yj=LogModSwitch[a1,Mod[bj PowerMod[aj,-xj,p],p],p,q,opts];
		If[yj==Null,Return[Null,Block]];
		xj=Mod[xj+q^(j-1) yj,q^j];
	,{j,2,k}];
	xj
];

(* Pohlig-Hellman reduction for coprime factos using CRT *)

LogMod[a_Integer,b_Integer,p_Integer,opts:OptionsPattern[LogMod]]:=Block[{m,mf,aj,bj,pj,kj,l,ll={},mm={}},
	If[GCD[a,p]>1,Return[Message[LogMod::noinv,a,p],Block]];
	m=MultiplicativeOrder[a,p];
	If[OptionValue[Verbose]>0,Print["LogMod: ord(a)=",m]];
	If[PowerMod[b,m,p]!=1,Return[Message[LogMod::nosol,a,b,p];Null,Block]];
	mf=FactorInteger[m];
	If[Length[mf]==1,Return[LogModPH[a,b,p,mf[[1,1]],mf[[1,2]],opts]]];
	Do[
		{pj,kj}=mf[[j]];
		aj=PowerMod[a,m/pj^kj,p];
		bj=PowerMod[b,m/pj^kj,p];
		l=LogModPH[aj,bj,p,pj,kj,opts];
		If[l==Null,Return[Null,Block]];
		mm=Append[mm,pj^kj];
		ll=Append[ll,l];
	,{j,Length[mf]}];
	ChineseRemainder[ll,mm]
];

End[]


Protect[{LogMod,LogModRho,LogModBSGS,LogModIC}];

EndPackage[]
