From 7894c218f592267c856a136ca2a422401a062fd9 Mon Sep 17 00:00:00 2001
From: Max Horn <max@quendi.de>
Date: Thu, 26 Sep 2024 23:11:07 +0200
Subject: [PATCH] Introduce `flintify` helper for "optimal" dispatch on integer
 and rational inputs (#1867)

---
 src/flint/FlintTypes.jl | 63 +++++++++++++++++++++++++++++++++++++++++
 src/flint/fmpq_mpoly.jl | 41 +++++++--------------------
 src/flint/fmpz_mpoly.jl | 17 +----------
 src/flint/fmpz_poly.jl  | 28 ++++++------------
 4 files changed, 82 insertions(+), 67 deletions(-)

diff --git a/src/flint/FlintTypes.jl b/src/flint/FlintTypes.jl
index 4051fd740..d9b03b320 100644
--- a/src/flint/FlintTypes.jl
+++ b/src/flint/FlintTypes.jl
@@ -1362,6 +1362,8 @@ mutable struct ZZMPolyRingElem <: MPolyRingElem{ZZRingElem}
     z.parent = ctx
     return z
   end
+
+  ZZMPolyRingElem(ctx::ZZMPolyRing, a::Integer) = ZZMPolyRingElem(ctx, flintify(a))
 end
 
 function _fmpz_mpoly_clear_fn(a::ZZMPolyRingElem)
@@ -6453,8 +6455,69 @@ end
 #
 ################################################################################
 
+"""
+    IntegerUnion = Union{Integer, ZZRingElem}
+
+The `IntegerUnion` type union allows convenient and compact declaration
+of methods that accept both Julia and Nemo integers.
+"""
 const IntegerUnion = Union{Integer, ZZRingElem}
 
+"""
+    RationalUnion = Union{Integer, ZZRingElem, Rational, QQFieldElem}
+
+The `RationalUnion` type union allows convenient and compact declaration
+of methods that accept both Julia and Nemo integers or rationals.
+"""
+const RationalUnion = Union{Integer, ZZRingElem, Rational, QQFieldElem}
+
+"""
+    flintify(x::RationalUnion)
+
+Return either an `Int`, `ZZRingElem` or `QQFieldElem` equal to `x`.
+
+This internal helper allow us to cleanly and compactly implement efficient
+dispatch for arithmetics that involve native Nemo objects plus a Julia
+integer.
+
+Indeed, many internal arithmetics functions in FLINT have optimize variants
+for the case when one operand is an `ZZRingElem` or an `Int` (sometimes also
+`UInt` is supported). E.g. there are special methods for adding one of these
+to a `ZZRingPolyElem`.
+
+In order to handling adding an arbitrary Julia integer to a `ZZRingPolyElem`,
+further dispatch is needed. The easiest is to provide a method
+
+    +(a::ZZRingPolyElem, b::Integer) = a + ZZ(b)
+
+However this is inefficient when `b` is e.g. an `UInt16`. So we could
+do this (at least on a 64 bit machine):
+
+    +(a::ZZRingPolyElem, b::Integer) = a + ZZ(b)
+    +(a::ZZRingPolyElem, b::{Int64,Int32,Int16,Int8,UInt32,UInt16,UInt8}) = a + Int(b)
+
+Doing this repeatedly is cumbersome and error prone. This can be avoided by
+using `flintify`, which allows us to write
+
+    +(a::ZZRingPolyElem, b::Integer) = a + flintify(b)
+
+to get optimal dispatch.
+
+This also works for Nemo types that also have special handlers for `UInt`,
+as their method for `b::UInt` takes precedence over the fallback method.
+"""
+flintify(x::ZZRingElem) = x
+flintify(x::QQFieldElem) = x
+flintify(x::Int) = x
+flintify(x::Integer) = ZZRingElem(x)::ZZRingElem
+flintify(x::Rational) = QQFieldElem(x)::QQFieldElem
+@static if Int === Int64
+  flintify(x::Union{Int64,Int32,Int16,Int8,UInt32,UInt16,UInt8}) = Int(x)
+else
+  flintify(x::Union{Int32,Int16,Int8,UInt16,UInt8}) = Int(x)
+end
+
+
 const ZmodNFmpzPolyRing = Union{ZZModPolyRing, FpPolyRing}
 
 const Zmodn_poly = Union{zzModPolyRingElem, fpPolyRingElem}
diff --git a/src/flint/fmpq_mpoly.jl b/src/flint/fmpq_mpoly.jl
index e926383bf..24cbea21b 100644
--- a/src/flint/fmpq_mpoly.jl
+++ b/src/flint/fmpq_mpoly.jl
@@ -369,13 +369,13 @@ for (jT, cN, cT) in ((QQFieldElem, :fmpq, Ref{QQFieldElem}), (ZZRingElem, :fmpz,
   end
 end
 
-+(a::QQMPolyRingElem, b::Integer) = a + ZZRingElem(b)
++(a::QQMPolyRingElem, b::Integer) = a + flintify(b)
 
 +(a::Integer, b::QQMPolyRingElem) = b + a
 
--(a::QQMPolyRingElem, b::Integer) = a - ZZRingElem(b)
+-(a::QQMPolyRingElem, b::Integer) = a - flintify(b)
 
--(a::Integer, b::QQMPolyRingElem) = -(b - a)
+-(a::Integer, b::QQMPolyRingElem) = neg!(b - a)
 
 +(a::QQMPolyRingElem, b::Rational{<:Integer}) = a + QQFieldElem(b)
 
@@ -383,9 +383,9 @@ end
 
 -(a::QQMPolyRingElem, b::Rational{<:Integer}) = a - QQFieldElem(b)
 
--(a::Rational{<:Integer}, b::QQMPolyRingElem) = -(b - a)
+-(a::Rational{<:Integer}, b::QQMPolyRingElem) = neg!(b - a)
 
-*(a::QQMPolyRingElem, b::Integer) = a * ZZRingElem(b)
+*(a::QQMPolyRingElem, b::Integer) = a * flintify(b)
 
 *(a::Integer, b::QQMPolyRingElem) = b * a
 
@@ -393,11 +393,11 @@ end
 
 *(a::Rational{<:Integer}, b::QQMPolyRingElem) = b * a
 
-divexact(a::QQMPolyRingElem, b::Integer; check::Bool=true) = divexact(a, ZZRingElem(b); check=check)
+divexact(a::QQMPolyRingElem, b::Integer; check::Bool=true) = divexact(a, flintify(b); check=check)
 
 divexact(a::QQMPolyRingElem, b::Rational{<:Integer}; check::Bool=true) = divexact(a, QQFieldElem(b); check=check)
 
-//(a::QQMPolyRingElem, b::Integer) = //(a, ZZRingElem(b))
+//(a::QQMPolyRingElem, b::Integer) = //(a, flintify(b))
 
 //(a::QQMPolyRingElem, b::Rational{<:Integer}) = //(a, QQFieldElem(b))
 
@@ -578,7 +578,7 @@ end
 
 ==(a::Int, b::QQMPolyRingElem) = b == a
 
-==(a::QQMPolyRingElem, b::Integer) = a == ZZRingElem(b)
+==(a::QQMPolyRingElem, b::Integer) = a == flintify(b)
 
 ==(a::Integer, b::QQMPolyRingElem) = b == a
 
@@ -1186,33 +1186,12 @@ function (R::QQMPolyRing)()
   return z
 end
 
-function (R::QQMPolyRing)(b::QQFieldElem)
+function (R::QQMPolyRing)(b::RationalUnion)
   z = QQMPolyRingElem(R, b)
   return z
 end
 
-function (R::QQMPolyRing)(b::ZZRingElem)
-  z = QQMPolyRingElem(R, b)
-  return z
-end
-
-function (R::QQMPolyRing)(b::Int)
-  z = QQMPolyRingElem(R, b)
-  return z
-end
-
-function (R::QQMPolyRing)(b::UInt)
-  z = QQMPolyRingElem(R, b)
-  return z
-end
-
-function (R::QQMPolyRing)(b::Integer)
-  return R(ZZRingElem(b))
-end
-
-function (R::QQMPolyRing)(b::Rational{<:Integer})
-  return R(QQFieldElem(b))
-end
+QQMPolyRingElem(ctx::QQMPolyRing, a::RationalUnion) = QQMPolyRingElem(ctx, flintify(a))
 
 function (R::QQMPolyRing)(a::QQMPolyRingElem)
   parent(a) != R && error("Unable to coerce polynomial")
diff --git a/src/flint/fmpz_mpoly.jl b/src/flint/fmpz_mpoly.jl
index b00790e6a..3ab6dfeff 100644
--- a/src/flint/fmpz_mpoly.jl
+++ b/src/flint/fmpz_mpoly.jl
@@ -1066,26 +1066,11 @@ function (R::ZZMPolyRing)()
   return z
 end
 
-function (R::ZZMPolyRing)(b::ZZRingElem)
+function (R::ZZMPolyRing)(b::IntegerUnion)
   z = ZZMPolyRingElem(R, b)
   return z
 end
 
-function (R::ZZMPolyRing)(b::Int)
-  z = ZZMPolyRingElem(R, b)
-  return z
-end
-
-function (R::ZZMPolyRing)(b::UInt)
-  z = ZZMPolyRingElem(R, b)
-  return z
-end
-
-function (R::ZZMPolyRing)(b::Integer)
-  return R(ZZRingElem(b))
-end
-
-
 function (R::ZZMPolyRing)(a::ZZMPolyRingElem)
   parent(a) != R && error("Unable to coerce polynomial")
   return a
diff --git a/src/flint/fmpz_poly.jl b/src/flint/fmpz_poly.jl
index d8bf129fc..c1ccf333e 100644
--- a/src/flint/fmpz_poly.jl
+++ b/src/flint/fmpz_poly.jl
@@ -225,17 +225,17 @@ end
 
 *(x::ZZPolyRingElem, y::ZZRingElem) = y*x
 
-+(x::Integer, y::ZZPolyRingElem) = y + ZZRingElem(x)
++(x::Integer, y::ZZPolyRingElem) = y + flintify(x)
 
--(x::Integer, y::ZZPolyRingElem) = ZZRingElem(x) - y
+-(x::Integer, y::ZZPolyRingElem) = flintify(x) - y
 
-*(x::Integer, y::ZZPolyRingElem) = ZZRingElem(x)*y
+*(x::Integer, y::ZZPolyRingElem) = flintify(x)*y
 
-+(x::ZZPolyRingElem, y::Integer) = x + ZZRingElem(y)
++(x::ZZPolyRingElem, y::Integer) = x + flintify(y)
 
--(x::ZZPolyRingElem, y::Integer) = x - ZZRingElem(y)
+-(x::ZZPolyRingElem, y::Integer) = x - flintify(y)
 
-*(x::ZZPolyRingElem, y::Integer) = ZZRingElem(y)*x
+*(x::ZZPolyRingElem, y::Integer) = flintify(y)*x
 
 ###############################################################################
 #
@@ -942,20 +942,8 @@ function (a::ZZPolyRing)()
   return z
 end
 
-function (a::ZZPolyRing)(b::Int)
-  z = ZZPolyRingElem(b)
-  z.parent = a
-  return z
-end
-
-function (a::ZZPolyRing)(b::Integer)
-  z = ZZPolyRingElem(ZZRingElem(b))
-  z.parent = a
-  return z
-end
-
-function (a::ZZPolyRing)(b::ZZRingElem)
-  z = ZZPolyRingElem(b)
+function (a::ZZPolyRing)(b::IntegerUnion)
+  z = ZZPolyRingElem(flintify(b))
   z.parent = a
   return z
 end