use relation graph for gravity calculations

this should improve performance a bit, since we don't call
gravityForces for every object anymore.

this change was achieved by making the relation graph directed,
while enforcing that every (directed) edge is accompanied by one
in the opposite direction.

6 files changed, 124 insertions, 91 deletions

diff --git a/grav2ty.cabal b/grav2ty.cabal
index da3f3cc..fc0ef22 100644
--- a/grav2ty.cabal
+++ b/grav2ty.cabal
@@ -21,7 +21,7 @@ extra-source-files:  CHANGELOG.md
 library grav2ty-lib
   exposed-modules:     Grav2ty.Simulation
                      , Grav2ty.Control
-                     , Grav2ty.Util.UGraph
+                     , Grav2ty.Util.RelGraph
   -- other-extensions:
   build-depends:       base ^>=4.12.0.0
                      , containers ^>=0.6.0.1
diff --git a/lib/Grav2ty/Control.hs b/lib/Grav2ty/Control.hs
index ea90af5..0b39667 100644
--- a/lib/Grav2ty/Control.hs
+++ b/lib/Grav2ty/Control.hs
@@ -13,7 +13,7 @@ module Grav2ty.Control
   ) where
 
 import Grav2ty.Simulation
-import Grav2ty.Util.UGraph
+import Grav2ty.Util.RelGraph
 
 import Control.Lens
 import Data.Foldable
@@ -93,7 +93,7 @@ applyControls cs obj@Dynamic {} =
 
 type ExtractFunction a b = Object a -> Maybe (State a b -> State a b)
 
-updateState :: (RealFloat a, Ord a) => a -> ExtractFunction a b
+updateState :: (Show a, RealFloat a, Ord a) => a -> ExtractFunction a b
                 -> State a b -> State a b
 updateState t extract state =
   over (control.ctrlTick) (+ 1)
@@ -102,12 +102,13 @@ updateState t extract state =
   where oldWorld = state^.world
         (newWorld, updateState') = foldl' updateAndExtract (S.empty, Nothing) oldWorld
         updateAndExtract acc@(seq, f) x =
-          if isDynamic x && (anyU relColl x objectRel == Just True)
+          if isDynamic x && (anyFrom _relColl x objectRel == Just True)
              then acc
              else (updateObject' x >< seq, chainFun (extract x) f)
         chainFun x f = if isJust f then (.) <$> x <*> f else x
         objectRel = objectRelGraph oldWorld
+        getForce obj = foldlFrom' (\f r -> f + _relForce r) (V2 0 0) obj objectRel
         scaledT = state^.control^.ctrlTimeScale * t
         updateObject' obj =
-          fmap (updateObject scaledT (gravitationForces oldWorld obj))
+          fmap (updateObject scaledT (getForce obj))
           . applyControls (state^.control) $ obj
diff --git a/lib/Grav2ty/Simulation.hs b/lib/Grav2ty/Simulation.hs
index 7a7af42..65bbf94 100644
--- a/lib/Grav2ty/Simulation.hs
+++ b/lib/Grav2ty/Simulation.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE TemplateHaskell #-}
 module Grav2ty.Simulation
   (
   -- * Objects
@@ -7,7 +8,7 @@ module Grav2ty.Simulation
   , isDynamic
   , World (..)
   , updateObject
-  , gravitationForces
+  , gravitationForce
   , ObjRel (..)
   , objectRelGraph
   -- * Hitboxes
@@ -22,7 +23,7 @@ module Grav2ty.Simulation
   , rotateV2
   ) where
 
-import Grav2ty.Util.UGraph
+import Grav2ty.Util.RelGraph
 
 import Control.Lens
 import Data.Complex
@@ -217,11 +218,19 @@ gravitationForces world obj = foldl' calcSum (pure 0) world
                              else force
 
 data ObjRel a = ObjRel
-  { relColl :: Bool
+  { _relColl  :: Bool -- ^ Wether the two 'Object's collide
+  , _relForce :: V2 a -- ^ The gravitation force between them
   } deriving (Show, Eq, Ord)
 
-objectRelGraph :: (RealFloat a, Ord a) => World a -> UGraph (Object a) (ObjRel a)
-objectRelGraph = insertSeq (\a b -> ObjRel $ objectCollision a b) emptyU
+makeLenses 'ObjRel
+
+objectRelGraph :: (RealFloat a, Ord a) => World a -> RelGraph (Object a) (ObjRel a)
+objectRelGraph = insertSeq rel emptyRel
+  where rel a b = let res = ObjRel (objectCollision a b) (gravity a b)
+                   in (res, over relForce negated res)
+        gravity a b = if separated a b
+                         then gravitationForce a b
+                         else pure 0
 
 updateObject :: Fractional a =>  a -> V2 a -> Object a -> Object a
 updateObject _ _ obj@Static {} = obj
diff --git a/lib/Grav2ty/Util/RelGraph.hs b/lib/Grav2ty/Util/RelGraph.hs
new file mode 100644
index 0000000..589fb91
--- /dev/null
+++ b/lib/Grav2ty/Util/RelGraph.hs
@@ -0,0 +1,94 @@
+module Grav2ty.Util.RelGraph
+  (
+  -- * Interface
+    RelGraph ()
+  , emptyRel
+  , insertRel
+  , insertRelNoOv
+  , insertSeq
+  , lookupRel
+  , anyFrom
+  , foldlFrom'
+  -- * Props
+  , prop_relCorrectness
+  , prop_insertLookup
+  , prop_insertLookupNoOv
+  ) where
+
+import Data.Foldable
+import Data.Map.Strict (Map (..))
+import Data.Maybe
+import Data.Sequence (Seq (..), (<|))
+import qualified Data.Map.Strict as M
+
+-- | Representation of a directed Relation Graph.
+--   Every Edge between two vertices @a@ has a connected value @v@.
+--   Every Edge must have an Edge in the opposite direction.
+--   A Vertice must not be connected to itself.
+--
+--   These “laws” are enforced by the functions of this interface.
+newtype RelGraph a v = RelGraph { unRelGraph :: Map a (Map a v) }
+  deriving (Show, Read, Eq)
+
+-- | Relation Graph without any Edges or Vertices.
+emptyRel :: RelGraph a v
+emptyRel = RelGraph M.empty
+
+insertInternal :: Ord a => (a -> v -> Map a v -> Map a v) -> a -> a -> v -> v
+               -> RelGraph a v -> RelGraph a v
+insertInternal f x y v v' =
+  if x == y
+     then id
+     else RelGraph . M.alter (alter v' x) y
+            . M.alter (alter v y) x . unRelGraph
+  where alter v i Nothing  = Just $ M.singleton i v
+        alter v i (Just m) = Just $ f i v m
+
+-- | Insert an edge with two connected values (for either direction)
+--   into an 'RelGraph'. Will ignore identity relations.
+insertRel :: Ord a => a -> a -> v -> v -> RelGraph a v -> RelGraph a v
+insertRel = insertInternal M.insert
+
+-- | Like 'insertRel', but won't overwrite any existing values.
+insertRelNoOv :: Ord a => a -> a -> v -> v -> RelGraph a v -> RelGraph a v
+insertRelNoOv = insertInternal (M.insertWith (\_ old -> old))
+
+-- | Takes a 'Seq' of Vertices and a function that returns the relations for
+--   the associated edge and inserts them into a 'RelGraph'
+insertSeq :: Ord a => (a -> a -> (v, v)) -> RelGraph a v -> Seq a -> RelGraph a v
+insertSeq f g seq = ins seq g
+  where ins (x :<| s) acc = ins s (foldl' (folder x) acc s)
+        ins mempty acc = acc
+        folder x g el = let (v, v') = f x el
+                       in insertRelNoOv x el v v' g
+
+-- | Lookup the Relation between two given Vertices.
+lookupRel :: Ord a => a -> a -> RelGraph a v -> Maybe v
+lookupRel x y = (>>= (M.lookup y)) . M.lookup x . unRelGraph
+
+-- | Wether any of the Edges from a given Vertex satisfy
+--   the given condition.
+anyFrom :: Ord a => (v -> Bool) -> a -> RelGraph a v -> Maybe Bool
+anyFrom f x = fmap (foldl (\b x -> b || f x) False) . M.lookup x . unRelGraph
+
+-- | Strict foldl over the Edges from a given Vertex
+foldlFrom' :: Ord a => (b -> v -> b) -> b -> a -> RelGraph a v -> b
+foldlFrom' f res x =
+  foldl' f res . fromMaybe M.empty . M.lookup x . unRelGraph
+
+prop_relCorrectness :: Seq Integer -> Bool
+prop_relCorrectness seq = and . map cond $ allCombinations g
+  where g = insertSeq (\x y -> let r = x * y in (r, negate r)) emptyRel seq
+        cond (x, y) = lookupRel x y g == fmap negate (lookupRel y x g)
+
+prop_insertLookupNoOv :: (Ord a, Eq v) => a -> a -> v ->  Bool
+prop_insertLookupNoOv x y v =
+  x == y || Just v == lookupRel x y (insertRelNoOv x y v v emptyRel)
+
+prop_insertLookup :: (Ord a, Eq v) => a -> a -> v -> RelGraph a v -> Bool
+prop_insertLookup x y v g =
+  x == y || Just v == lookupRel x y (insertRel x y v v g)
+
+allCombinations :: RelGraph a v -> [(a, a)]
+allCombinations (RelGraph m) = foldl' (\l k -> map (\x -> (k, x)) keys ++ l) [] keys
+  where keys = M.keys m
diff --git a/lib/Grav2ty/Util/UGraph.hs b/lib/Grav2ty/Util/UGraph.hs
deleted file mode 100644
index ad75520..0000000
--- a/lib/Grav2ty/Util/UGraph.hs
+++ /dev/null
@@ -1,71 +0,0 @@
-module Grav2ty.Util.UGraph
-  (
-  -- * Interface
-    UGraph ()
-  , emptyU
-  , insertU
-  , insertSeq
-  , lookupU
-  , anyU
-  -- * Props
-  , prop_undirected
-  , prop_undirected'
-  , prop_insertLookup
-  , prop_insertLookupNoOv
-  ) where
-
-import Data.Foldable
-import Data.Map.Strict (Map (..))
-import Data.Sequence (Seq (..), (<|))
-import qualified Data.Map.Strict as M
-
--- | Representation of an undirected Graph.
-newtype UGraph a v = UGraph { unUGraph :: Map a (Map a v) } deriving (Show, Read, Eq)
-
--- | Undirected Graph without any Edges or Vertices.
-emptyU :: UGraph a v
-emptyU = UGraph M.empty
-
-alterOuter :: Ord a => (a -> v -> Map a v -> Map a v) -> v -> a
-           -> Maybe (Map a v) -> Maybe (Map a v)
-alterOuter f v i Nothing  = Just $ M.singleton i v
-alterOuter f v i (Just m) = Just $ f i v m
-
--- | Insert an edge with a connected value into an 'UGraph'
-insertU :: Ord a => a -> a -> v -> UGraph a v -> UGraph a v
-insertU x y v = UGraph . M.alter (alter x) y
-  . M.alter (alter y) x . unUGraph
-  where alter = alterOuter M.insert v
-
-insertUNoOv :: Ord a => a -> a -> v -> UGraph a v -> UGraph a v
-insertUNoOv x y v = UGraph . M.alter (alter y) x
-  . M.alter (alter x) y . unUGraph
-  where alter = alterOuter (M.insertWith (\_ old -> old)) v
-
-insertSeq :: Ord a => (a -> a -> v) -> UGraph a v -> Seq a -> UGraph a v
-insertSeq f g seq = ins seq g
-  where ins (x :<| s) acc =
-          ins s (foldl' (\g el -> insertUNoOv x el (f x el) g) acc s)
-        ins mempty acc = acc
-
-lookupU :: Ord a => a -> a -> UGraph a v -> Maybe v
-lookupU x y = (>>= (M.lookup y)) . M.lookup x . unUGraph
-
-anyU :: Ord a => (v -> Bool) -> a -> UGraph a v -> Maybe Bool
-anyU f x = fmap (foldl (\b x -> b || f x) False) . M.lookup x . unUGraph
-
-prop_undirected' :: (Ord a, Eq v) => UGraph a v -> a -> a -> Bool
-prop_undirected' g x y = lookupU x y g == lookupU y x g
-
-prop_undirected :: (Ord a, Eq v) => UGraph a v -> Bool
-prop_undirected g = and . map (uncurry (prop_undirected' g)) $ allCombinations g
-
-prop_insertLookupNoOv :: (Ord a, Eq v) => a -> a -> v ->  Bool
-prop_insertLookupNoOv x y v = Just v == lookupU x y (insertUNoOv x y v emptyU)
-
-prop_insertLookup :: (Ord a, Eq v) => a -> a -> v -> UGraph a v ->  Bool
-prop_insertLookup x y v g = Just v == lookupU x y (insertU x y v g)
-
-allCombinations :: UGraph a v -> [(a, a)]
-allCombinations (UGraph m) = foldl' (\l k -> map (\x -> (k, x)) keys ++ l) [] keys
-  where keys = M.keys m
diff --git a/test/Main.hs b/test/Main.hs
index c85665b..2d52c51 100644
--- a/test/Main.hs
+++ b/test/Main.hs
@@ -3,33 +3,33 @@ module Main where
 import Test.Tasty as QC
 import Test.Tasty.QuickCheck as QC
 
-import Grav2ty.Util.UGraph
+import Grav2ty.Util.RelGraph
 
 times :: Monad m => Int -> a -> (a -> m a) -> m a
 times i res f = if i > 0
                    then f res >>= flip (times (i - 1)) f
                    else pure res
 
-instance (Ord a, Arbitrary a, Arbitrary v) => Arbitrary (UGraph a v) where
+instance (Ord a, Arbitrary a, Arbitrary v) => Arbitrary (RelGraph a v) where
   arbitrary = do
-    let g = emptyU
+    let g = emptyRel
     c <- arbitrary `suchThat` (>= 0)
     times c g $ (\g -> do
       x <- arbitrary
       y <- arbitrary
       v <- arbitrary
-      pure $ insertU x y v g)
+      v' <- arbitrary
+      pure $ insertRel x y v v' g)
 
-uGraph :: TestTree
-uGraph = testGroup "Grav2ty.Util.UGraph"
-  [ QC.testProperty "Check undirectedness" (prop_undirected :: UGraph Int Bool -> Bool)
-  , QC.testProperty "Check undirectedness (with vertices not in the graph)" (prop_undirected' :: UGraph Char Int -> Char -> Char -> Bool)
-  , QC.testProperty "Check insertion and lookup correctness" (prop_insertLookup :: Integer -> Integer ->  String -> UGraph Integer String -> Bool)
+relGraph :: TestTree
+relGraph = testGroup "Grav2ty.Util.UGraph"
+  [ QC.testProperty "Check relational properties of insertSeq" prop_relCorrectness
+  , QC.testProperty "Check insertion and lookup correctness" (prop_insertLookup :: Integer -> Integer ->  String -> RelGraph Integer String -> Bool)
   , QC.testProperty "Check insertion and lookup correctness w/o overwrite" (prop_insertLookupNoOv :: Integer -> Integer ->  String -> Bool)
   ]
 
 libTests :: TestTree
-libTests = testGroup "Grav2ty Library Tests" [ uGraph ]
+libTests = testGroup "Grav2ty Library Tests" [ relGraph ]
 
 main :: IO ()
 main = defaultMain libTests