1 files changed, 18 insertions, 13 deletions

diff --git a/edit-lens/src/Data/String/DFST/Lens.lhs b/edit-lens/src/Data/String/DFST/Lens.lhs
index 9def01b..8fd748a 100644
--- a/edit-lens/src/Data/String/DFST/Lens.lhs
+++ b/edit-lens/src/Data/String/DFST/Lens.lhs
@@ -1,5 +1,6 @@
 \begin{code}
 {-# LANGUAGE ScopedTypeVariables
+           , TemplateHaskell
 #-}
 
 module Data.String.DFST.Lens
@@ -14,6 +15,8 @@ module Data.String.DFST.Lens
 
 import Data.String.DFST
 import Control.Lens.Edit
+import Control.Lens
+import Control.Lens.TH
 import Control.Edit
 
 import Numeric.Natural
@@ -27,20 +30,24 @@ import Data.Algorithm.Diff (Diff, getDiff)
 import qualified Data.Algorithm.Diff as Diff
 
 import Data.Monoid
+import Data.Function (on)
 
 
-data StringEdit = Insert Natural Char
-                | Delete Natural
+data StringEdit = Insert { _sePos :: Natural, _seInsertion :: Char }
+                | Delete { _sePos :: Natural }
   deriving (Eq, Ord, Show, Read)
 
-sePos :: StringEdit -> Natural
-sePos (Insert pos _) = pos
-sePos (Delete pos  ) = pos
+makeLenses ''StringEdit
 
 data StringEdits = StringEdits (Seq StringEdit)
                  | SEFail
   deriving (Eq, Ord, Show, Read)
 
+makePrisms ''StringEdits
+
+stringEdits :: Traversal' StringEdits StringEdit
+stringEdits = _StringEdits . traverse
+
 insert :: Natural -> Char -> StringEdits
 insert n c = StringEdits .  Seq.singleton $ Insert n c
   
@@ -121,9 +128,9 @@ dfstLens dfst@DFST{..} = EditLens ground propR propL
 
     propR :: (DFSTComplement state, StringEdits) -> (DFSTComplement state, StringEdits)
     propR (c, SEFail) = (c, SEFail)
-    propR (c, StringEdits (es :|> e)) = (c', es' <> shiftBy (length $ pOutput []) (strDiff sOutput sOutput'))
+    propR (c, StringEdits (es :|> e)) = (c', es' <> es'')
       where
-        (cSuffix, cPrefix) = Comp.splitAt (Comp.length c - (fromEnum $ sePos e)) c
+        (cSuffix, cPrefix) = Comp.splitAt (Comp.length c - (e ^. sePos . from enum)) c
         cSuffix'
           | Delete _       <- e = Comp.take (pred $ Comp.length cSuffix) cSuffix -- TODO unsafe
           | Insert _ nChar <- e = cSuffix <> Comp.singleton (DFSTAction $ \x -> runDFST' dfst x (pure nChar) id)
@@ -131,18 +138,16 @@ dfstLens dfst@DFST{..} = EditLens ground propR propL
         (_, sOutput ) = runDFSTAction (Comp.composed cSuffix ) pState
         (_, sOutput') = runDFSTAction (Comp.composed cSuffix') pState
         (c', es') = propR (cSuffix' <> cPrefix, StringEdits es)
-        shiftBy (toEnum -> n) (StringEdits es) = StringEdits $ shiftBy' n <$> es
-        shiftBy' n' (Insert n c) = Insert (n + n') c
-        shiftBy' n' (Delete n) = Delete (n + n')
+        es'' = (strDiff `on` ($ "")) sOutput sOutput' & stringEdits . sePos . from enum +~ (length $ pOutput [])
     propR (c, mempty) = (c, mempty)
         
 
     propL :: (DFSTComplement state, StringEdits) -> (DFSTComplement state, StringEdits)
     propL = undefined
 
-strDiff :: (String -> String) -> (String -> String) -> StringEdits
--- ^ @strDiff a b@ calculates a set of edits, which, when applied to @a@, produce @b@, where both @a@ and @b@ are given as diff-lists of chars
-strDiff a b = snd . foldr toEdit (0, mempty) $ getDiff (a []) (b [])
+strDiff :: String -> String -> StringEdits
+-- ^ @strDiff a b@ calculates a set of edits, which, when applied to @a@, produce @b@
+strDiff a b = snd . foldr toEdit (0, mempty) $ getDiff a b
   where
     toEdit :: Diff Char -> (Natural, StringEdits) -> (Natural, StringEdits)
     toEdit (Diff.Both _ _) (n, es) = (succ n, es)