content: -- |
  -- Module:      Data.Poly.Internal.Dense.GcdDomain
  -- Copyright:   (c) 2019 Andrew Lelechenko
  -- Licence:     BSD3
  -- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
  --
  -- 'GcdDomain' instance with a 'GcdDomain' constraint on the coefficient type.
  --
  
  {-# LANGUAGE FlexibleInstances          #-}
  {-# LANGUAGE MultiWayIf                 #-}
  {-# LANGUAGE ScopedTypeVariables        #-}
  {-# LANGUAGE TypeFamilies               #-}
  
  {-# OPTIONS_GHC -fno-warn-orphans #-}
  
  module Data.Poly.Internal.Dense.GcdDomain
    () where
  
  import Prelude hiding (gcd, lcm, (^))
  import Control.Exception
  import Control.Monad
  import Control.Monad.ST
  import Data.Euclidean
  import Data.Maybe
  import Data.Semiring (Semiring(..), Ring(), isZero, minus)
  import qualified Data.Vector.Generic as G
  import qualified Data.Vector.Generic.Mutable as MG
  
  import Data.Poly.Internal.Dense
  
  -- | @since 0.3.0.0
  instance (Eq a, Ring a, GcdDomain a, G.Vector v a) => GcdDomain (Poly v a) where
    divide (Poly xs) (Poly ys) =
      toPoly' <$> quotient xs ys
    {-# INLINABLE divide #-}
  
    gcd (Poly xs) (Poly ys)
      | G.null xs = Poly ys
      | G.null ys = Poly xs
      | G.length xs == 1 = Poly $ G.singleton $ G.foldl' gcd (G.unsafeHead xs) ys
      | G.length ys == 1 = Poly $ G.singleton $ G.foldl' gcd (G.unsafeHead ys) xs
      | otherwise = toPoly' $ gcdNonEmpty xs ys
    {-# INLINE gcd #-}
  
    lcm x@(Poly xs) y@(Poly ys)
      | G.null xs || G.null ys = zero
      | otherwise = (x `divide'` gcd x y) `times` y
    {-# INLINABLE lcm #-}
  
    coprime x y = isJust (one `divide` gcd x y)
    {-# INLINABLE coprime #-}
  
  gcdNonEmpty
    :: (Eq a, Ring a, GcdDomain a, G.Vector v a)
    => v a
    -> v a
    -> v a
  gcdNonEmpty xs ys = runST $ do
      let x = G.foldl1' gcd xs
          y = G.foldl1' gcd ys
          xy = x `gcd` y
      xs' <- G.thaw xs
      ys' <- G.thaw ys
      zs' <- gcdM xs' ys'
  
      let lenZs = MG.length zs'
          go acc 0 = pure acc
          go acc n = do
            t <- MG.unsafeRead zs' (n - 1)
            go (acc `gcd` t) (n - 1)
      a <- MG.unsafeRead zs' (lenZs - 1)
      z <- go a (lenZs - 1)
  
      forM_ [0 .. lenZs - 1] $ \i ->
        MG.unsafeModify zs'((`times` xy) . (`divide'` z)) i
  
      G.unsafeFreeze zs'
  {-# INLINABLE gcdNonEmpty #-}
  
  gcdM
    :: (Eq a, Ring a, GcdDomain a, G.Vector v a)
    => G.Mutable v s a
    -> G.Mutable v s a
    -> ST s (G.Mutable v s a)
  gcdM xs ys
    | MG.null xs = pure ys
    | MG.null ys = pure xs
    | otherwise = do
    let lenXs = MG.length xs
        lenYs = MG.length ys
    xLast <- MG.unsafeRead xs (lenXs - 1)
    yLast <- MG.unsafeRead ys (lenYs - 1)
    let z  = xLast `lcm` yLast
        zx = z `divide'` xLast
        zy = z `divide'` yLast
  
    if
      | lenYs <= lenXs
      , Just xy <- xLast `divide` yLast -> do
        forM_ [0 .. lenYs - 1] $ \i -> do
          y <- MG.unsafeRead ys i
          when (y /= zero) $
            MG.unsafeModify
              xs
              (\x -> x `minus` y `times` xy)
              (i + lenXs - lenYs)
        xs' <- dropWhileEndM isZero xs
        gcdM xs' ys
      | lenXs <= lenYs
      , Just yx <- yLast `divide` xLast -> do
        forM_ [0 .. lenXs - 1] $ \i -> do
          x <- MG.unsafeRead xs i
          when (x /= zero) $
            MG.unsafeModify
              ys
              (\y -> y `minus` x `times` yx)
              (i + lenYs - lenXs)
        ys' <- dropWhileEndM isZero ys
        gcdM xs ys'
      | lenYs <= lenXs -> do
        forM_ [0 .. lenYs - 1] $ \i -> do
          y <- MG.unsafeRead ys i
          MG.unsafeModify
            xs
            (\x -> x `times` zx `minus` y `times` zy)
            (i + lenXs - lenYs)
        forM_ [0 .. lenXs - lenYs - 1] $
          MG.unsafeModify xs (`times` zx)
        xs' <- dropWhileEndM isZero xs
        gcdM xs' ys
      | otherwise -> do
        forM_ [0 .. lenXs - 1] $ \i -> do
          x <- MG.unsafeRead xs i
          MG.unsafeModify
            ys
            (\y -> y `times` zy `minus` x `times` zx)
            (i + lenYs - lenXs)
        forM_ [0 .. lenYs - lenXs - 1] $
          MG.unsafeModify ys (`times` zy)
        ys' <- dropWhileEndM isZero ys
        gcdM xs ys'
  {-# INLINABLE gcdM #-}
  
  divide' :: GcdDomain a => a -> a -> a
  divide' = (fromMaybe (error "gcd: violated internal invariant") .) . divide
  
  isZeroM
    :: (Eq a, Semiring a, G.Vector v a)
    => G.Mutable v s a
    -> ST s Bool
  isZeroM xs = go (MG.length xs)
    where
      go 0 = pure True
      go n = do
        x <- MG.unsafeRead xs (n - 1)
        if x == zero then go (n - 1) else pure False
  {-# INLINE isZeroM #-}
  
  quotient
    :: (Eq a, Ring a, GcdDomain a, G.Vector v a)
    => v a
    -> v a
    -> Maybe (v a)
  quotient xs ys
    | G.null ys = throw DivideByZero
    | G.null xs = Just xs
    | G.length xs < G.length ys = Nothing
    | otherwise = runST $ do
      let lenXs = G.length xs
          lenYs = G.length ys
          lenQs = lenXs - lenYs + 1
      qs <- MG.unsafeNew lenQs
      rs <- MG.unsafeNew lenXs
      G.unsafeCopy rs xs
  
      let go i
            | i < 0 = do
              b <- isZeroM rs
              if b
                then Just <$> G.unsafeFreeze qs
                else pure Nothing
            | otherwise = do
              r <- MG.unsafeRead rs (lenYs - 1 + i)
              case r `divide` G.unsafeLast ys of
                Nothing -> pure Nothing
                Just q -> do
                  MG.unsafeWrite qs i q
                  forM_ [0 .. lenYs - 1] $ \k ->
                    MG.unsafeModify
                      rs
                      (\c -> c `minus` q `times` G.unsafeIndex ys k)
                      (i + k)
                  go (i - 1)
  
      go (lenQs - 1)
  {-# INLINABLE quotient #-}
  