{-# LANGUAGE ScopedTypeVariables #-}
module RefDiffBy (diffBy) where

import Data.TreeDiff.List (Edit (..))

import qualified Data.Primitive as P

diffBy :: forall a. (a -> a -> Bool) -> [a] -> [a] -> [Edit a]
diffBy eq xs' ys' = reverse (getCell (lcs xn yn))
  where
    xn = length xs'
    yn = length ys'

    xs = P.arrayFromListN xn xs'
    ys = P.arrayFromListN yn ys'

    memo :: P.Array (Cell [Edit a])
    memo = P.arrayFromListN ((xn + 1) * (yn + 1))
        [ impl xi yi
        | xi <- [0 .. xn]
        , yi <- [0 .. yn]
        ]

    lcs :: Int -> Int -> Cell [Edit a]
    lcs xi yi = P.indexArray memo (yi + xi * (yn + 1))

    impl :: Int -> Int -> Cell [Edit a]
    impl 0 0 = Cell 0 []
    impl 0 m = case lcs 0 (m - 1) of
        Cell w edit -> Cell (w + 1) (Ins (P.indexArray ys (m - 1)) : edit)
    impl n 0 = case lcs (n - 1) 0 of
        Cell w edit -> Cell (w + 1) (Del (P.indexArray xs (n - 1)) : edit)

    impl n m = bestOfThree
        edit
        (bimap (+1) (Ins y :) (lcs n (m - 1)))
        (bimap (+1) (Del x :) (lcs (n - 1) m))
      where
        x = P.indexArray xs (n - 1)
        y = P.indexArray ys (m - 1)

        edit
            | eq x y    = bimap id   (Cpy x :)   (lcs (n - 1) (m - 1))
            | otherwise = bimap (+1) (Swp x y :) (lcs (n - 1) (m - 1))

data Cell a = Cell !Int !a

getCell :: Cell a -> a
getCell (Cell _ x) = x

bestOfThree :: Cell a -> Cell a -> Cell a -> Cell a
bestOfThree a@(Cell i _x) b@(Cell j _y) c@(Cell k _z)
    | i <= j
    = if i <= k then a else c

    | otherwise
    = if j <= k then b else c

bimap :: (Int -> Int) -> (a -> b) -> Cell a -> Cell b
bimap f g (Cell i x) = Cell (f i) (g x)