never executed always true always false
    1 {-|
    2 Module: Flaw.Visual.Geometry.Simplification
    3 Description: Geometry simplification algorithm.
    4 License: MIT
    5 -}
    6 
    7 {-# LANGUAGE FlexibleContexts, ViewPatterns #-}
    8 
    9 module Flaw.Visual.Geometry.Simplification
   10   ( simplifyGeometry
   11   ) where
   12 
   13 import Control.Monad
   14 import Control.Monad.ST
   15 import qualified Data.Map.Strict as M
   16 import Data.STRef
   17 import qualified Data.Vector.Generic as VG
   18 import qualified Data.Vector.Generic.Mutable as VGM
   19 import qualified Data.Vector.Storable as VS
   20 import qualified Data.Vector.Storable.Mutable as VSM
   21 import qualified Data.Vector.Unboxed.Mutable as VUM
   22 import Data.Word
   23 import Foreign.Ptr
   24 import Foreign.Storable
   25 
   26 import Flaw.Math
   27 
   28 data PairKey = PairKey {-# UNPACK #-} !Word32 {-# UNPACK #-} !Word32 deriving (Eq, Ord)
   29 
   30 data Pair = Pair
   31   { pair_vertex0 :: {-# UNPACK #-} !Word32
   32   , pair_vertex1 :: {-# UNPACK #-} !Word32
   33   , pair_optimalVertex :: {-# UNPACK #-} !Float3
   34   , pair_cost :: {-# UNPACK #-} !Float
   35   }
   36 
   37 instance Storable Pair where
   38   sizeOf _ = 24
   39   alignment _ = 4
   40   peek ptr = Pair
   41     <$> peek (castPtr ptr)
   42     <*> peek (castPtr ptr `plusPtr` 4)
   43     <*> peek (castPtr ptr `plusPtr` 8)
   44     <*> peek (castPtr ptr `plusPtr` 20)
   45   poke ptr (Pair v0 v1 ov c) = do
   46     poke (castPtr ptr) v0
   47     poke (castPtr ptr `plusPtr` 4) v1
   48     poke (castPtr ptr `plusPtr` 8) ov
   49     poke (castPtr ptr `plusPtr` 20) c
   50 
   51 -- | Simplify geometry.
   52 -- Algorithm: http://cseweb.ucsd.edu/~ravir/190/2016/garland97.pdf
   53 simplifyGeometry :: Int -> VS.Vector Float3 -> VS.Vector Word32 -> (VS.Vector Float3, VS.Vector Word32)
   54 simplifyGeometry iterationsCount vertices indices = runST $ do
   55   let
   56     verticesCount = VG.length vertices
   57     indicesCount = VG.length indices
   58     trianglesCount = indicesCount `quot` 3
   59 
   60   -- allocate memory
   61   vertexPositions <- VS.thaw vertices
   62   vertexParents <- VUM.new verticesCount
   63   let
   64     f i = when (i < verticesCount) $ do
   65       VGM.unsafeWrite vertexParents i i
   66       f $ i + 1
   67     in f 0
   68   vertexRanks <- VUM.replicate verticesCount (1 :: Word32)
   69   vertexQuadrics <- VSM.replicate verticesCount 0
   70   pairHeap <- VSM.new indicesCount
   71   pairHeapSizeRef <- newSTRef 0
   72   pairHeapIndexByKeyRef <- newSTRef M.empty
   73 
   74   let
   75     calculatePair (PairKey v0 v1) = do
   76       quadric0 <- VGM.unsafeRead vertexQuadrics $ fromIntegral v0
   77       quadric1 <- VGM.unsafeRead vertexQuadrics $ fromIntegral v1
   78       let
   79         quadricSum@(Float4x4
   80           q11 q12 q13 q14
   81           q21 q22 q23 q24
   82           q31 q32 q33 q34
   83           _q41 _q42 _q43 _q44
   84           ) = quadric0 + quadric1
   85         q_12_12 = q11 * q22 - q21 * q12
   86         q_12_13 = q11 * q23 - q21 * q13
   87         q_12_14 = q11 * q24 - q21 * q14
   88         q_12_23 = q12 * q23 - q22 * q13
   89         q_12_24 = q12 * q24 - q22 * q14
   90         q_12_34 = q13 * q24 - q23 * q14
   91         q_123_123 = q_12_12 * q33 - q_12_13 * q32 + q_12_23 * q31
   92         q_123_124 = q_12_12 * q34 - q_12_14 * q32 + q_12_24 * q31
   93         q_123_134 = q_12_13 * q34 - q_12_14 * q33 + q_12_34 * q31
   94         q_123_234 = q_12_23 * q34 - q_12_24 * q33 + q_12_34 * q32
   95         -- q_1234_1234 = q_123_123 * q44 - q_123_124 * q43 + q_123_134 * q42 - q_123_234 * q41
   96         aff :: Float3 -> Float4
   97         aff (Float3 x y z) = Float4 x y z 1
   98         cost :: Float4x4 -> Float3 -> Float
   99         cost q (aff -> p) = p `dot` (q `mul` p)
  100       pv0 <- VGM.unsafeRead vertexPositions $ fromIntegral v0
  101       pv1 <- VGM.unsafeRead vertexPositions $ fromIntegral v1
  102       let
  103         (oc, ov) =
  104           -- try global minimum
  105           if abs q_123_123 > 1e-4 then let
  106             v = Float3
  107               (negate $ q_123_234 / q_123_123)
  108               (q_123_134 / q_123_123)
  109               (negate $ q_123_124 / q_123_123)
  110             in (cost quadricSum v, v)
  111           -- try minimum along the edge
  112           else let
  113             v10 = aff pv1 - aff pv0
  114             u = quadricSum `mul` v10
  115             h = u `dot` v10
  116             in if abs h > 1e-4 then let
  117               hinv = 1 / h
  118               v = pv0 * vecFromScalar ((u `dot` aff pv1) * hinv) - pv1 * vecFromScalar ((u `dot` aff pv0) * hinv)
  119               in (cost quadricSum v, v)
  120             -- try minimum from ends and the middle
  121             else let
  122               pm = (pv0 + pv1) * 0.5
  123               in min (cost quadricSum pm, pm) $ min (cost quadricSum pv0, pv0) (cost quadricSum pv1, pv1)
  124 
  125       return Pair
  126         { pair_vertex0 = v0
  127         , pair_vertex1 = v1
  128         , pair_optimalVertex = ov
  129         , pair_cost = oc - cost quadric0 pv0 - cost quadric1 pv1
  130         }
  131 
  132   -- calculate initial vertex quadrics
  133   let
  134     f i = when (i < trianglesCount) $ do
  135       let
  136         i1 = fromIntegral $ indices VG.! (i * 3)
  137         i2 = fromIntegral $ indices VG.! (i * 3 + 1)
  138         i3 = fromIntegral $ indices VG.! (i * 3 + 2)
  139         p1 = vertices VG.! i1
  140         p2 = vertices VG.! i2
  141         p3 = vertices VG.! i3
  142         normal@(Float3 a b c) = normalize $ cross (p2 - p1) (p3 - p1)
  143         d = negate $ dot normal p1
  144         aa = a * a
  145         ab = a * b
  146         ac = a * c
  147         ad = a * d
  148         bb = b * b
  149         bc = b * c
  150         bd = b * d
  151         cc = c * c
  152         cd = c * d
  153         dd = d * d
  154         q = Float4x4
  155           aa ab ac ad
  156           ab bb bc bd
  157           ac bc cc cd
  158           ad bd cd dd
  159       VGM.unsafeModify vertexQuadrics (+ q) i1
  160       VGM.unsafeModify vertexQuadrics (+ q) i2
  161       VGM.unsafeModify vertexQuadrics (+ q) i3
  162       f $ i + 1
  163     in f 0
  164 
  165   -- add initial pairs
  166   let
  167     f i = when (i < trianglesCount) $ do
  168       let
  169         i1 = indices VG.! (i * 3)
  170         i2 = indices VG.! (i * 3 + 1)
  171         i3 = indices VG.! (i * 3 + 2)
  172         addPair v0 v1 = unless (v0 == v1) $ do
  173           let key = PairKey (min v0 v1) (max v0 v1)
  174           pairHeapIndexByKey <- readSTRef pairHeapIndexByKeyRef
  175           unless (M.member key pairHeapIndexByKey) $ do
  176             pairHeapSize <- readSTRef pairHeapSizeRef
  177             VGM.unsafeWrite pairHeap pairHeapSize =<< calculatePair key
  178             writeSTRef pairHeapSizeRef $! pairHeapSize + 1
  179             writeSTRef pairHeapIndexByKeyRef $! M.insert key pairHeapSize pairHeapIndexByKey
  180       addPair i1 i2
  181       addPair i2 i3
  182       addPair i1 i3
  183       f $ i + 1
  184     in f 0
  185 
  186   -- add reverse pairs in index too
  187   do
  188     pairs <- readSTRef pairHeapIndexByKeyRef
  189     writeSTRef pairHeapIndexByKeyRef $! foldr (\(PairKey v0 v1, p) -> M.insert (PairKey v1 v0) p) pairs (M.toList pairs)
  190 
  191   -- find open edges and add additional costs for vertices
  192   do
  193     pairHeapIndexByKey <- readSTRef pairHeapIndexByKeyRef
  194     pairHeapSize <- readSTRef pairHeapSizeRef
  195     -- count triangles for every edge
  196     pairTriangleCounts <- VUM.replicate pairHeapSize (0 :: Word32)
  197     let
  198       f i = when (i < trianglesCount) $ do
  199         let
  200           i1 = indices VG.! (i * 3)
  201           i2 = indices VG.! (i * 3 + 1)
  202           i3 = indices VG.! (i * 3 + 2)
  203           accountForPair a b = case M.lookup (PairKey a b) pairHeapIndexByKey of
  204             Just h -> VGM.unsafeModify pairTriangleCounts (+ 1) h
  205             Nothing -> return ()
  206         accountForPair i1 i2
  207         accountForPair i2 i3
  208         accountForPair i1 i3
  209         f $ i + 1
  210       in f 0
  211     -- for every open edge add penalty to its vertices for moving from the edge
  212     let
  213       f i = when (i < pairHeapSize) $ do
  214         triangleCount <- VGM.unsafeRead pairTriangleCounts i
  215         when (triangleCount == 1) $ do
  216           Pair
  217             { pair_vertex0 = v0
  218             , pair_vertex1 = v1
  219             } <- VGM.unsafeRead pairHeap i
  220           pv0 <- VGM.unsafeRead vertexPositions $ fromIntegral v0
  221           pv1 <- VGM.unsafeRead vertexPositions $ fromIntegral v1
  222           when (norm2 (pv1 - pv0) > 1e-8) $ do
  223             let
  224               r@(Float3 rx ry rz) = normalize $ pv1 - pv0
  225               rxx = rx * rx
  226               rxy = rx * ry
  227               rxz = rx * rz
  228               ryy = ry * ry
  229               ryz = ry * rz
  230               rzz = rz * rz
  231               a = Float3 (rxx - 1) rxy rxz
  232               b = Float3 rxy (ryy - 1) ryz
  233               c = Float3 rxz ryz (rzz - 1)
  234               d = pv0 - r * vecFromScalar (dot r pv0)
  235               aa = dot a a
  236               ab = dot a b
  237               ac = dot a c
  238               ad = dot a d
  239               bb = dot b b
  240               bc = dot b c
  241               bd = dot b d
  242               cc = dot c c
  243               cd = dot c d
  244               dd = dot d d
  245               q = Float4x4
  246                 aa ab ac ad
  247                 ab bb bc bd
  248                 ac bc cc cd
  249                 ad bd cd dd
  250             VGM.unsafeModify vertexQuadrics (+ q) $ fromIntegral v0
  251             VGM.unsafeModify vertexQuadrics (+ q) $ fromIntegral v1
  252         f $ i + 1
  253       in f 0
  254 
  255   -- vertex disjoint-set functions
  256 
  257   let
  258     -- get vertex parent
  259     vertexParent a = do
  260       p <- VGM.unsafeRead vertexParents a
  261       if p == a then return p else do
  262         pp <- vertexParent p
  263         VGM.unsafeWrite vertexParents a pp
  264         return pp
  265     -- union vertices
  266     unionVertices a b = do
  267       pa <- vertexParent a
  268       pb <- vertexParent b
  269       if pa == pb then return pa else do
  270         ra <- VGM.unsafeRead vertexRanks pa
  271         rb <- VGM.unsafeRead vertexRanks pb
  272         if ra > rb then do
  273           VGM.unsafeWrite vertexParents pb pa
  274           VGM.unsafeWrite vertexRanks pa $ ra + rb
  275           return pa
  276         else do
  277           VGM.unsafeWrite vertexParents pa pb
  278           VGM.unsafeWrite vertexRanks pb $ ra + rb
  279           return pb
  280 
  281   -- heap functions
  282 
  283   -- swap elements
  284     heapSwap a pa@Pair
  285       { pair_vertex0 = a0
  286       , pair_vertex1 = a1
  287       } b pb@Pair
  288       { pair_vertex0 = b0
  289       , pair_vertex1 = b1
  290       } = do
  291       VGM.unsafeWrite pairHeap b pa
  292       VGM.unsafeWrite pairHeap a pb
  293       modifySTRef' pairHeapIndexByKeyRef
  294         $ M.insert (PairKey a0 a1) b
  295         . M.insert (PairKey a1 a0) b
  296         . M.insert (PairKey b0 b1) a
  297         . M.insert (PairKey b1 b0) a
  298 
  299     -- sift down
  300     heapSiftDown i = do
  301       pairHeapSize <- readSTRef pairHeapSizeRef
  302       p <- VGM.unsafeRead pairHeap i
  303       let
  304         l = i * 2 + 1
  305         r = i * 2 + 2
  306         swap m mp = do
  307           heapSwap i p m mp
  308           heapSiftDown m
  309       lp <- if l < pairHeapSize then VGM.unsafeRead pairHeap l else return p
  310       rp <- if r < pairHeapSize then VGM.unsafeRead pairHeap r else return p
  311       if pair_cost lp < pair_cost p && pair_cost lp <= pair_cost rp then swap l lp
  312       else if pair_cost rp < pair_cost p && pair_cost rp <= pair_cost lp then swap r rp
  313       else return ()
  314 
  315     -- sift up
  316     heapSiftUp i = when (i > 0) $ do
  317       p <- VGM.unsafeRead pairHeap i
  318       let m = (i - 1) `quot` 2
  319       mp <- VGM.unsafeRead pairHeap m
  320       when (pair_cost p < pair_cost mp) $ do
  321         heapSwap i p m mp
  322         heapSiftUp m
  323 
  324     -- update
  325     heapUpdate i =
  326       if i > 0 then do
  327         p <- VGM.unsafeRead pairHeap i
  328         let m = (i - 1) `quot` 2
  329         mp <- VGM.unsafeRead pairHeap m
  330         case compare (pair_cost p) (pair_cost mp) of
  331           LT -> heapSiftUp i
  332           GT -> heapSiftDown i
  333           EQ -> return ()
  334       else heapSiftDown i
  335 
  336     -- delete
  337     heapDelete i = do
  338       pairHeapSize <- readSTRef pairHeapSizeRef
  339       let m = pairHeapSize - 1
  340       writeSTRef pairHeapSizeRef $! m
  341       Pair
  342         { pair_vertex0 = v0
  343         , pair_vertex1 = v1
  344         } <- VGM.unsafeRead pairHeap i
  345       if i < m then do
  346         mp@Pair
  347           { pair_vertex0 = mv0
  348           , pair_vertex1 = mv1
  349           } <- VGM.unsafeRead pairHeap m
  350         VGM.unsafeWrite pairHeap i mp
  351         modifySTRef' pairHeapIndexByKeyRef
  352           $ M.insert (PairKey mv0 mv1) i
  353           . M.insert (PairKey mv1 mv0) i
  354           . M.delete (PairKey v0 v1)
  355           . M.delete (PairKey v1 v0)
  356         heapUpdate i
  357       else modifySTRef' pairHeapIndexByKeyRef
  358         $ M.delete (PairKey v0 v1)
  359         . M.delete (PairKey v1 v0)
  360 
  361   -- make actual heap from pairs
  362   do
  363     pairHeapSize <- readSTRef pairHeapSizeRef
  364     let
  365       f i = when (i >= 0) $ do
  366         heapSiftDown i
  367         f $ i - 1
  368       in f $ pairHeapSize - 1
  369 
  370   -- contraction step
  371   let
  372     contraction = do
  373       -- get minimal cost pair
  374       Pair
  375         { pair_vertex0 = pv0
  376         , pair_vertex1 = pv1
  377         , pair_optimalVertex = vo
  378         } <- VGM.unsafeRead pairHeap 0
  379       -- union vertices
  380       v <- unionVertices (fromIntegral pv0) (fromIntegral pv1)
  381       -- update position
  382       VGM.unsafeWrite vertexPositions v vo
  383       -- delete minimal cost pair
  384       heapDelete 0
  385       -- figure out what vertex got replaced
  386       let
  387         (v0, v1) = case (fromIntegral v == pv0, fromIntegral v == pv1) of
  388           (True, False) -> (pv0, pv1)
  389           (False, True) -> (pv1, pv0)
  390           _ -> error $ show ("impossible pair", v, pv0, pv1)
  391       -- sum up quadrics
  392       v1q <- VGM.unsafeRead vertexQuadrics (fromIntegral v1)
  393       VGM.unsafeModify vertexQuadrics (+ v1q) (fromIntegral v0)
  394       -- update all pairs with v1 to use v0
  395       pairsWithV1
  396         <-  M.takeWhileAntitone (\(PairKey a _) -> a == v1)
  397         .   M.dropWhileAntitone (\(PairKey a _) -> a < v1)
  398         <$> readSTRef pairHeapIndexByKeyRef
  399       forM_ (M.keys pairsWithV1) $ \oldKey@(PairKey _v1 b) -> unless (v0 == b) $ do
  400         let newKey = PairKey v0 b
  401         pairHeapIndex <- readSTRef pairHeapIndexByKeyRef
  402         Just h <- return $ M.lookup oldKey pairHeapIndex
  403         -- if contraction made double edge, remove it
  404         if M.member newKey pairHeapIndex then heapDelete h
  405         -- else fix edge to point to v0
  406         else do
  407           writeSTRef pairHeapIndexByKeyRef $
  408             ( M.insert newKey h
  409             . M.insert (PairKey b v0) h
  410             . M.delete oldKey
  411             . M.delete (PairKey b v1)
  412             ) pairHeapIndex
  413           VGM.unsafeModify pairHeap (\p -> p
  414             { pair_vertex0 = min b v0
  415             , pair_vertex1 = max b v0
  416             }) h
  417           heapUpdate h
  418       -- update all pairs using v0
  419       pairsWithV0
  420         <-  M.takeWhileAntitone (\(PairKey a _) -> a == v0)
  421         .   M.dropWhileAntitone (\(PairKey a _) -> a < v0)
  422         <$> readSTRef pairHeapIndexByKeyRef
  423       forM_ (M.keys pairsWithV0) $ \(PairKey _v0 b) -> do
  424         Just h <- M.lookup (PairKey v0 b) <$> readSTRef pairHeapIndexByKeyRef
  425         VGM.unsafeWrite pairHeap h =<< calculatePair (PairKey (min v0 b) (max v0 b))
  426         heapUpdate h
  427 
  428   -- perform contractions
  429   replicateM_ iterationsCount contraction
  430 
  431   -- filter out indices with degenerate triangles
  432   newIndices <- VSM.new indicesCount
  433   end <- let
  434     f p i =
  435       if i < trianglesCount then do
  436         i1 <- vertexParent $ fromIntegral $ indices VG.! (i * 3)
  437         i2 <- vertexParent $ fromIntegral $ indices VG.! (i * 3 + 1)
  438         i3 <- vertexParent $ fromIntegral $ indices VG.! (i * 3 + 2)
  439         if i1 == i2 || i1 == i3 || i2 == i3 then f p (i + 1) else do
  440           VSM.unsafeWrite newIndices p $ fromIntegral i1
  441           VSM.unsafeWrite newIndices (p + 1) $ fromIntegral i2
  442           VSM.unsafeWrite newIndices (p + 2) $ fromIntegral i3
  443           f (p + 3) (i + 1)
  444       else return p
  445     in f 0 0
  446 
  447   resultVertices <- VG.unsafeFreeze vertexPositions
  448   resultIndices <- VG.unsafeFreeze $ VGM.slice 0 end newIndices
  449 
  450   return (resultVertices, resultIndices)