```-- | Call graphs and related concepts, more or less as defined in
--     \"A Predicative Analysis of Structural Recursion\" by
--     Andreas Abel and Thorsten Altenkirch.

-- Originally copied from Agda1 sources.

module Agda.Termination.CallGraph
( -- * Structural orderings
Order(..)
, (.*.)
, supremum
-- * Call matrices
, Index
, CallMatrix(..)
, (>*<)
, callMatrixInvariant
-- * Calls
, Call(..)
, callInvariant
-- * Call graphs
, CallGraph
, callGraphInvariant
, fromList
, toList
, empty
, union
, insert
, complete
, showBehaviour
-- * Tests
, Agda.Termination.CallGraph.tests
) where

import Agda.Utils.QuickCheck
import Agda.Utils.Function
import Agda.Utils.List
import Agda.Utils.TestHelpers
import Agda.Termination.Matrix as Matrix
import Agda.Termination.Semiring (Semiring)
import qualified Agda.Termination.Semiring as Semiring
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Map (Map, (!))
import qualified Data.Map as Map
import Data.List hiding (union, insert)
import Data.Monoid
import Data.Array (elems)
import Data.Function

------------------------------------------------------------------------
-- Structural orderings

-- | The order called R in the paper referred to above. Note that
-- @'Unknown' '<=' 'Le' '<=' 'Lt'@.
--
--
-- TODO: document orders which are call-matrices themselves.
data Order
= Lt | Le | Unknown | Mat (Matrix Integer Order)
deriving (Eq,Ord)

instance Show Order where
show Lt      = "<"
show Le      = "="
show Unknown = "?"
show (Mat m) = "Mat " ++ show m

--instance Ord Order where
--    max = maxO

instance Arbitrary Order where
arbitrary = elements [Lt, Le, Unknown]

instance CoArbitrary Order where
coarbitrary Lt      = variant 0
coarbitrary Le      = variant 1
coarbitrary Unknown = variant 2
coarbitrary (Mat m) = variant 3

-- | Multiplication of 'Order's. (Corresponds to sequential
-- composition.)

(.*.) :: Order -> Order -> Order
Lt      .*. Unknown   = Unknown
Lt      .*. (Mat m)   = Lt .*. (collapse m)
Lt      .*. _         = Lt
Le      .*. o         = o
Unknown .*. _         = Unknown
(Mat m1) .*. (Mat m2) = if (okM m1 m2) then
Mat \$ mul orderSemiring m1 m2
else
(collapse m1) .*. (collapse m2)
(Mat m) .*. Le        = Mat m
(Mat m) .*. Unknown   = Unknown
(Mat m) .*. Lt        = (collapse m) .*. Lt

collapse :: Matrix Integer Order -> Order
collapse m = foldl (.*.) Le (Data.Array.elems \$ diagonal m)

okM :: Matrix Integer Order -> Matrix Integer Order -> Bool
okM m1 m2 = (rows \$ size m2) == (cols \$ size m1)

-- | The supremum of a (possibly empty) list of 'Order's.

supremum :: [Order] -> Order
supremum = foldr maxO Unknown

maxO :: Order -> Order -> Order
maxO o1 o2 = case (o1,o2) of
(_,Lt) -> Lt
(Lt,_) -> Lt
(Unknown,_) -> o2
(_,Unknown) -> o1
(Mat m1, Mat m2) -> Mat (Matrix.zipWith maxO m1 m2)
(Mat m,_) -> maxO (collapse m) o2
(_,Mat m) ->  maxO o1 (collapse m)
(Le,Le) -> Le

-- | The infimum of a (possibly empty) list of 'Order's.

-- infimum :: [Order] -> Order
-- infimum = foldr min Lt -- DELETE ?

-- | @('Order', 'max', '.*.')@ forms a semiring, with 'Unknown' as zero
-- and 'Le' as one.

orderSemiring :: Semiring Order
orderSemiring =
, Semiring.mul = (.*.)
, Semiring.zero = Unknown
, Semiring.one = Le
}

prop_orderSemiring = Semiring.semiringInvariant orderSemiring

------------------------------------------------------------------------
-- Call matrices

-- | Call matrix indices.

type Index = Integer

-- | Call matrices. Note the call matrix invariant
-- ('callMatrixInvariant').

newtype CallMatrix = CallMatrix { mat :: Matrix Index Order }
deriving (Eq, Ord, Show)

instance Arbitrary CallMatrix where
arbitrary = do
sz <- arbitrary
callMatrix sz

instance CoArbitrary CallMatrix where
coarbitrary (CallMatrix m) = coarbitrary m

prop_Arbitrary_CallMatrix = callMatrixInvariant

-- | Generates a call matrix of the given size.

callMatrix :: Size Index -> Gen CallMatrix
callMatrix sz = do
m <- matrixUsingRowGen sz rowGen
return \$ CallMatrix { mat = m }
where
rowGen :: Index -> Gen [Order]
rowGen 0 = return []
rowGen n = do
x <- arbitrary
i <- choose (0, n - 1)
return \$ genericReplicate i Unknown ++ [x] ++
genericReplicate (n - 1 - i) Unknown

prop_callMatrix sz =
forAll (callMatrix sz) \$ \cm ->
callMatrixInvariant cm
&&
size (mat cm) == sz

-- | In a call matrix at most one element per row may be different
-- from 'Unknown'.

callMatrixInvariant :: CallMatrix -> Bool
callMatrixInvariant cm =
matrixInvariant m &&
all ((<= 1) . length . filter (/= Unknown)) (toLists m)
where m = mat cm

-- | Call matrix multiplication.
--
-- Precondition: see 'mul'.

(<*>) :: CallMatrix -> CallMatrix -> CallMatrix
cm1 <*> cm2 =
CallMatrix { mat = mul orderSemiring (mat cm1) (mat cm2) }

prop_cmMul sz =
forAll natural \$ \c2 ->
forAll (callMatrix sz) \$ \cm1 ->
forAll (callMatrix \$ Size { rows = cols sz, cols = c2 }) \$ \cm2 ->
callMatrixInvariant (cm1 <*> cm2)

------------------------------------------------------------------------
-- Calls

-- | This datatype encodes information about a single recursive
-- function application. The columns of the call matrix stand for
-- 'source' function arguments (patterns); the first argument has
-- index 0, the second 1, and so on. The rows of the matrix stand for
-- 'target' function arguments. Element @(i, j)@ in the matrix should
-- be computed as follows:
--
--   * 'Lt' (less than) if the @j@-th argument to the 'target'
--     function is structurally strictly smaller than the @i@-th
--     pattern.
--
--   * 'Le' (less than or equal) if the @j@-th argument to the
--     'target' function is structurally smaller than the @i@-th
--     pattern.
--
--   * 'Unknown' otherwise.
--
--   The structural ordering used is defined in the paper referred to
--   above.

data Call =
Call { source :: Index        -- ^ The function making the call.
, target :: Index        -- ^ The function being called.
, cm :: CallMatrix       -- ^ The call matrix describing the call.
}
deriving (Eq, Ord, Show)

instance Arbitrary Call where
arbitrary = do
[s, t]    <- vectorOf 2 arbitrary
cm        <- arbitrary
return (Call { source = s, target = t, cm = cm })

instance CoArbitrary Call where
coarbitrary (Call s t cm) =
coarbitrary s . coarbitrary t . coarbitrary cm

prop_Arbitrary_Call :: Call -> Bool
prop_Arbitrary_Call = callInvariant

-- | 'Call' invariant.

callInvariant :: Call -> Bool
callInvariant = callMatrixInvariant . cm

-- | 'Call' combination.
--
-- Precondition: see '<*>'; furthermore the 'source' of the first
-- argument should be equal to the 'target' of the second one.

(>*<) :: Call -> Call -> Call
c1 >*< c2 =
Call { source    = source c2
, target    = target c1
, cm        = cm c1 <*> cm c2
}

------------------------------------------------------------------------
-- Call graphs

-- | A call graph is a set of calls. Every call also has some
-- associated meta information, which should be 'Monoid'al so that the
-- meta information for different calls can be combined when the calls
-- are combined.

newtype CallGraph meta = CallGraph { cg :: Map Call meta }
deriving (Eq, Show)

-- | 'CallGraph' invariant.

callGraphInvariant :: CallGraph meta -> Bool
callGraphInvariant = all (callInvariant . fst) . toList

-- | Converts a call graph to a list of calls with associated meta
-- information.

toList :: CallGraph meta -> [(Call, meta)]
toList = Map.toList . cg

-- | Converts a list of calls with associated meta information to a
-- call graph.

fromList :: Monoid meta => [(Call, meta)] -> CallGraph meta
fromList = CallGraph . Map.fromListWith mappend

-- | Creates an empty call graph.

empty :: CallGraph meta
empty = CallGraph Map.empty

-- | Takes the union of two call graphs.

union :: Monoid meta
=> CallGraph meta -> CallGraph meta -> CallGraph meta
union cs1 cs2 = CallGraph \$ (Map.unionWith mappend `on` cg) cs1 cs2

-- | Inserts a call into a call graph.

insert :: Monoid meta
=> Call -> meta -> CallGraph meta -> CallGraph meta
insert c m = CallGraph . Map.insertWith mappend c m . cg

-- | Generates a call graph.

callGraph :: (Monoid meta, Arbitrary meta) => Gen (CallGraph meta)
callGraph = do
indices <- fmap nub arbitrary
n <- natural
let noMatrices | null indices = 0
| otherwise    = n `min` 3  -- Not too many.
fmap fromList \$ vectorOf noMatrices (matGen indices)
where
matGen indices = do
[s, t] <- vectorOf 2 (elements indices)
[c, r] <- vectorOf 2 (choose (0, 2))     -- Not too large.
m <- callMatrix (Size { rows = r, cols = c })
callId <- arbitrary
return (Call { source = s, target = t, cm = m }, callId)

prop_callGraph =
forAll (callGraph :: Gen (CallGraph [Integer])) \$ \cs ->
callGraphInvariant cs

-- | Call graph combination. (Application of '>*<' to all pairs @(c1,
-- c2)@ for which @'source' c1 = 'target' c2@.)
--
-- Precondition: see '<*>'.

combine
:: Monoid meta => CallGraph meta -> CallGraph meta -> CallGraph meta
combine s1 s2 = fromList \$
[ (c1 >*< c2, m1 `mappend` m2)
| (c1, m1) <- toList s1, (c2, m2) <- toList s2
, source c1 == target c2
]

-- | @'complete' cs@ completes the call graph @cs@. A call graph is
-- complete if it contains all indirect calls; if @f -> g@ and @g ->
-- h@ are present in the graph, then @f -> h@ should also be present.

complete :: Monoid meta => CallGraph meta -> CallGraph meta
complete cs = complete' safeCS
where
safeCS = ensureCompletePrecondition cs

complete' cs | cs' .==. cs = cs
| otherwise   = complete' cs'
where
cs' = cs `union` combine cs safeCS
(.==.) = ((==) `on` (Map.keys . cg))

prop_complete =
forAll (callGraph :: Gen (CallGraph [Integer])) \$ \cs ->
isComplete (complete cs)

-- | Returns 'True' iff the call graph is complete.

isComplete :: (Ord meta, Monoid meta) => CallGraph meta -> Bool
isComplete s = all (`Map.member` cg s) combinations
where
calls = toList s
combinations =
[ c2 >*< c1 | (c1, _) <- calls, (c2, _) <- calls
, target c1 == source c2 ]

-- | Checks whether every 'Index' used in the call graph corresponds
-- to a fixed number of arguments (i.e. rows\/columns).

completePrecondition :: CallGraph meta -> Bool
completePrecondition cs =
all (allEqual . map snd) \$
groupOn fst \$
concat [ [(source c, cols \$ size' c), (target c, rows \$ size' c)]
| (c, _) <- toList cs]
where
size' = size . mat . cm

-- | Returns a call graph padded with 'Unknown's in such a way that
-- 'completePrecondition' is satisfied.

ensureCompletePrecondition
:: Monoid meta => CallGraph meta -> CallGraph meta
ensureCompletePrecondition cs =
CallGraph \$ Map.mapKeysWith mappend pad \$ cg cs
where
-- The maximum number of arguments detected for every index.
noArgs :: Map Index Integer
noArgs = foldr (\c m -> insert (source c) (cols' c) \$
insert (target c) (rows' c) m)
Map.empty
(map fst \$ toList cs)
where insert = Map.insertWith max

pad c = c { cm = CallMatrix { mat = padRows \$ padCols \$ mat \$ cm c } }
where
padCols = iterate' ((noArgs ! source c) - cols' c)

padRows = iterate' ((noArgs ! target c) - rows' c)

cols'  = cols . size'
rows'  = rows . size'
size'  = size . mat . cm

prop_ensureCompletePrecondition =
forAll (callGraph :: Gen (CallGraph [Integer])) \$ \cs ->
let cs' = ensureCompletePrecondition cs in
completePrecondition cs'
&&
all callInvariant (map fst \$ toList cs')
&&
and [ or [ new .==. old | (old, _) <- toList cs ]
| (new, _) <- toList cs' ]
where
c1 .==. c2 = all (all (uncurry (==)))
((zipZip `on` (toLists . mat . cm)) c1 c2)

-- zipZip discards the new elements.
zipZip :: [[a]] -> [[b]] -> [[(a, b)]]
zipZip xs ys = map (uncurry zip) \$ zip xs ys

-- | Displays the recursion behaviour corresponding to a call graph.

showBehaviour :: Show meta => CallGraph meta -> String
showBehaviour = concatMap showCall . toList
where
showCall (c, meta) | source c /= target c = ""
| otherwise            = unlines
[ "Function:  " ++ show (source c)
, "Behaviour: " ++ show (elems \$ diagonal \$ mat \$ cm c)
, "Meta info: " ++ show meta
]

------------------------------------------------------------------------
-- All tests

tests :: IO Bool
tests = runTests "Agda.Termination.CallGraph"
[ quickCheck' prop_orderSemiring
, quickCheck' prop_Arbitrary_CallMatrix
, quickCheck' prop_callMatrix
, quickCheck' prop_cmMul
, quickCheck' prop_Arbitrary_Call
, quickCheck' prop_callGraph
, quickCheck' prop_complete
, quickCheck' prop_ensureCompletePrecondition
]
```