packages feed

tree-edit-distance-0.1.0.0: src/TreeEditDistance.hs

{-# LANGUAGE BangPatterns, OverloadedStrings, FlexibleContexts, DeriveGeneric, DeriveAnyClass, Strict, 
    StrictData, MultiParamTypeClasses, FlexibleInstances, InstanceSigs, TupleSections, MagicHash, UnboxedTuples #-}
{-|
Module      :  TreeEditDistance
Copyright   :  (C) 2024 Tillmann Vogt

License     :  BSD-style (see the file LICENSE)
Maintainer  :  Tillmann Vogt <tillk.vogt@gmail.com>
Stability   :  provisional
Portability :  POSIX
-}
module TreeEditDistance
    ( editDistance, zhangShasha, addPostOrderIndexing, addLeftMosts, leftMostNodes, keyRoots, dEdge, extractNodes, fromAdj, show',
      partIsomorphic, graphsAreIsomorphic,
      AGraph, Label(Label), EdgeLabel(EdgeLabel), EdgeAttribute
    ) where

import           Data.List (groupBy, sortBy, foldl')
import           Data.Map (Map)
import qualified Data.Map as Map
import           Data.Maybe(catMaybes, maybeToList)
import qualified Data.Vector.Unboxed as VU
import           Data.Word(Word32)
import           Graph.IntMap (Graph(..), Edge8(..), EdgeAttribute, ExtractNodeType, adjustNode, children, lookupNode)
import qualified Graph.IntMap as Graph
import Data.Array
-- import Debug.Trace

-- | If there is a faster/better algorithm, it will be put here
editDistance :: AGraph -> AGraph -> Int
editDistance tree0 tree1 = zhangShasha tree0 tree1

type AGraph = Graph Label EdgeLabel
data Label =
  Label { name :: String,
          indx :: Word32, -- a second index where by construction the left node is always smaller than the right node
          leftMost :: Word32
        } deriving (Show)
data EdgeLabel = EdgeLabel String deriving (Show)


-- Adds and index to every node int the order: left child, right child, parent, 
-- see https://youtu.be/hwiks-n7vso?t=410  . That way the index of the parent is
-- always bigger than its children and the index of the right node is always higher than the left node
addPostOrderIndexing :: Word32 -> (Word32, AGraph) -> (Word32, AGraph)
addPostOrderIndexing n (count, graph)
    | null cs = inc (count, graph)
    | otherwise = inc (foldr addPostOrderIndexing (count, graph) (reverse cs))
  where cs = VU.toList (children graph n (EdgeLabel ""))
        inc (c,g) = (c+1, adjustNode (\nl -> nl { indx = c + 1 }) n g)

-- Adds to every node the left most descendant leaf node
--           n
--         / | \
--        /  |  \
--       /   |   \
--      /    |    \
--     /     |     \
--    /      |      \
--   /       |       \
--  /        |        \
-- /         |         \   cs = children
--(head cs) (tail cs)  ..
--firstNode  rest      ..   -- applying addLeftMosts recursively
-- (leftMostNode, leftMostGraph) (_,restGraph)
--
addLeftMosts :: Word32 -> (Word32, AGraph) -> (Word32, AGraph)
addLeftMosts n (_, graph) -- Why _ ? node is not used when input
    | null cs = (n, adjust graph n n)
    | otherwise = (leftMostNode, adjust restGraph n leftMostNode)
  where cs = VU.toList (children graph n (EdgeLabel ""))
        adjust g node leftMn = adjustNode (\nl -> nl { leftMost = leftMn }) node g

        -- children are split into the first node and the rest
        (leftMostNode, leftMostGraph) -- first node
          | null cs = (0, adjust graph n n) -- Why 0? node is not used when tuple is input
          | otherwise = addLeftMosts (head cs) (0, graph)
        (_,restGraph) | null cs = (0, adjust graph n n) -- rest
                                                 -- Why 0? node is not used
                      | otherwise = foldr addLeftMosts (0, leftMostGraph) (tail cs)

-- walks through the tree in the same way as addLeftMosts
leftMostNodes :: Word32 -> AGraph -> [Word32]
leftMostNodes n graph
    | null cs = nList
    | otherwise = (concatMap (\c -> leftMostNodes c graph) cs) ++ nList -- nList at the end ~> post order
  where cs = VU.toList (children graph n (EdgeLabel ""))
        nList = maybeToList (fmap leftMost (lookupNode n graph))

keyRoots :: [Word32] -> [Word32]
keyRoots l = [fromIntegral i + 1 | (i, x) <- zip [0..] l, notElem x (drop (i + 1) l)]

type DistArray = Array (Word32, Word32) Int

-- A good explanation: https://www.youtube.com/watch?v=6Ur8B35xCj8
zhangShasha :: AGraph -> AGraph -> Int
zhangShasha tree0 tree1 = -- Debug.Trace.trace (" finalArray\n" ++ show' finalArray) $
                          finalArray ! (flen l0, flen l1)
  where
    lmGraph tree = snd (addLeftMosts root (addPostOrderIndexing root (0,tree)))
    root = 0

    lmGraph0 = lmGraph tree0
    lmGraph1 = lmGraph tree1

    l0 = leftMostNodes root lmGraph0
    l1 = leftMostNodes root lmGraph1

    li ngraph ls = catMaybes (map nList ls)
      where nList n = fmap indx (lookupNode n ngraph)

    l0i = li lmGraph0 l0
    l1i = li lmGraph1 l1

    roots = [(k0,k1) | k0 <- keyRoots l0, k1 <- keyRoots l1]

    finalArray :: DistArray
    finalArray = foldl' (treeDist tree0 tree1 l0i l1i) startArray roots

    startArray :: DistArray
    startArray = listArray ((0,0),(flen l0, flen l1)) (take (le0 * le1) (repeat 0))
      where le0 = (length l0) + 1
            le1 = (length l1) + 1

flen :: [Word32] -> Word32
flen = fromIntegral . length


treeDist :: AGraph -> AGraph -> [Word32] -> [Word32] -> DistArray -> (Word32, Word32) -> DistArray
treeDist tree0 tree1 l0 l1 ar (i, j) = tdAr // [((i,j), forestAr ! (i,j))]
  where ts = [(i1,j1) | i1 <- [(l0 !! (fromIntegral i -1))..i] ,
                        j1 <- [(l1 !! (fromIntegral j -1))..j]]
        (tdAr, forestAr) = foldl' forestDist (ar, startForestArray) ts

        startForestArray = listArray ((0,0),(i, j)) l
          where l = firstRow ++
                    (concat (take (fromIntegral i) (map firstColumn firstCol)))
        firstColumn c = c : (take (fromIntegral j) (repeat 0))

        firstRow = (take j1' (repeat 0)) ++ (take (fromIntegral j + 1 - j1') [1..])
        firstCol = (take (i1'-1) (repeat 0)) ++ (take (fromIntegral i - i1' + 1) ([1..]))

        i1' = fromIntegral (l0 !! (fromIntegral i - 1))
        j1' = fromIntegral (l1 !! (fromIntegral j - 1))

        forestDist :: (DistArray , DistArray) -> (Word32, Word32) -> (DistArray, DistArray)
        forestDist (td, fArr) (i1,j1)
            | l0 !! (fromIntegral i1 - 1) == l0 !! (fromIntegral i - 1) &&
              l1 !! (fromIntegral j1 - 1) == l1 !! (fromIntegral j - 1) =
                          (td // [((i1,j1), m)], fArr // [((i1,j1), m)])
            | otherwise = (td, fArr // [((i1,j1), minimum [fd0, fd1, fd3])])
          where (delete,insert,relabel) = (1,1,1)

                fd0 = (fArr ! (i_temp, j1)) + delete
                fd1 = (fArr ! (i1, j_temp)) + insert
                fd2 = (fArr ! (i_temp, j_temp)) + cost
                m = minimum [fd0, fd1, fd2]
                fd3 = (fArr ! (i_temp2, j_temp2)) + (td ! (i1,j1))

                i_temp = if l0 !! (fromIntegral i - 1) > i1 - 1 then 0 else i1 - 1
                j_temp = if l1 !! (fromIntegral j - 1) > j1 - 1 then 0 else j1 - 1

                i1_temp = l0 !! (fromIntegral i1 - 1) - 1
                j1_temp = l1 !! (fromIntegral j1 - 1) - 1

                i_temp2 = if l0 !! (fromIntegral i - 1) > i1_temp then 0 else i1_temp
                j_temp2 = if l1 !! (fromIntegral j - 1) > j1_temp then 0 else j1_temp

                cost | fmap name (lookupNode (i1-1) tree0) ==
                       fmap name (lookupNode (j1-1) tree1) = 0
                     | otherwise = relabel


---------------------------------------------------------------------------------------------------

instance EdgeAttribute EdgeLabel where
  fastEdgeAttr (EdgeLabel _) = 0
  edgeFromAttr = Map.fromList [(0, EdgeLabel "")]
  show_e _ = "Edge"
  bases _ = [Edge8 0]

instance Enum Label where
  toEnum n = Label (show n) 0 0 
  fromEnum (Label n _ _) = read n

instance ExtractNodeType Label where
    extractNodeType (Label _ _ _) = ""

-- | Extract nodes from both starting node and adjacent nodes
extractNodes :: [(Word32, [Word32], e)] -> Map.Map Word32 Label
extractNodes adj = Map.fromList (map labelNode nodes)
  where
    nodes = (map sel1 adj) ++ (concat (map sel2 adj))
    sel1 (x, _, _) = x
    sel2 (_, y, _) = y
    labelNode n = ( n, Label (show n) 0 0)

dEdge :: EdgeLabel
dEdge = EdgeLabel ""

fromAdj :: Map Word32 nl -> [(Word32, [Word32], EdgeLabel)] -> Graph nl EdgeLabel
fromAdj nodesMap adj = foldl (newNodes nodesMap) Graph.empty adj
  where
    newNodes :: Map Word32 nl -> Graph nl EdgeLabel -> (Word32, [Word32], EdgeLabel) -> Graph nl EdgeLabel
    newNodes nm g (n,ns,eLabel) = Graph.insertEdges (Just True) edges $
                                  Graph.insertNodes lookedUpNodes g
      where
        lookedUpNodes = catMaybes $ map addLabel (n:ns)
        addLabel n1 = fmap (\label -> (n1,label)) (Map.lookup n1 nm)
        edges = zip es edgeLbls
        es = zipWith (\n0 n1 -> (n0, n1)) (repeat n) ns
        edgeLbls = repeat eLabel

show' :: Array (Word32, Word32) Int -> String
show' a = concat (map ((++ "\n") . show . map snd) (groupBy groupx sorted))
  where sorted = sortBy (\((x,_),_) ((y,_),_) -> compare x y) (assocs a)
        groupx ((x0,_),_) ((x1,_),_) = x0 == x1


partIsomorphic :: (Word32, AGraph) -> (Word32, AGraph) -> ([(Word32, Word32)], [(Word32, Word32)])
partIsomorphic (f0, graph0) (f1, graph1)
  | graphsAreIsomorphic (graph0, graph1) (f0, f1) = ([(f0,f1)],[])
  | otherwise = ([],[(f0,f1)])

instance Eq Label where Label l0 i0 lm0 == Label l1 i1 lm1 = l0 == l1

graphsAreIsomorphic :: (AGraph, AGraph) -> (Word32, Word32) -> Bool
graphsAreIsomorphic (graph0, graph1) (f0, f1)
  | label0 /= label1 || length children0 /= length children1 = False
  | otherwise = and (map (graphsAreIsomorphic (graph0, graph1)) (zip children0 children1))
  | otherwise = False
  where
    label0 = Graph.lookupNode f0 graph0
    label1 = Graph.lookupNode f1 graph1
    children0 = VU.toList (Graph.children graph0 f0 (EdgeLabel ""))
    children1 = VU.toList (Graph.children graph1 f1 (EdgeLabel ""))