module Language.Haskell.Tools.Transform.RangeTemplateToSourceTemplate where
import Control.Monad.State
import Control.Reference ((^.), _1)
import Data.Map as Map
import Data.Ord (Ord(..), Ordering(..))
import Data.Set as Set
import FastString (mkFastString)
import Language.Haskell.Tools.AST
import Language.Haskell.Tools.Transform.RangeTemplate
import Language.Haskell.Tools.Transform.SourceTemplate
import SrcLoc
import StringBuffer (StringBuffer, nextChar, atEnd)
rangeToSource :: SourceInfoTraversal node => StringBuffer -> Ann node dom RngTemplateStage
-> Ann node dom SrcTemplateStage
rangeToSource srcInput tree = let locIndices = getLocIndices tree
srcMap = mapLocIndices srcInput locIndices
in applyFragments (Map.elems srcMap) tree
getLocIndices :: SourceInfoTraversal e => Ann e dom RngTemplateStage -> Set (RealSrcLoc, Int)
getLocIndices = snd . flip execState (0, Set.empty) .
sourceInfoTraverseDown (SourceInfoTrf
(\ni -> do { mapM_ (\el -> case getRangeElemSpan el of Just sp -> modify (insertElem sp); _ -> return ()) (ni ^. rngTemplateNodeElems); return ni })
(\ni -> do { mapM_ (modify . insertElem) (ni ^. rngTmpSeparators); return ni })
pure )
(return ()) (return ())
where insertElem sp (i,m) = (i+1, Set.insert (realSrcSpanEnd sp, i) m)
mapLocIndices :: Ord k => StringBuffer -> Set (RealSrcLoc, k) -> Map k String
mapLocIndices inp = (^. _1) . Set.foldl (\(new, str, pos) (sp, k) -> let (rem, val, newPos) = takeSpan str pos sp
in (Map.insert k (reverse val) new, rem, newPos))
(Map.empty, inp, mkRealSrcLoc (mkFastString "") 1 1)
where takeSpan :: StringBuffer -> RealSrcLoc -> RealSrcLoc -> (StringBuffer, String, RealSrcLoc)
takeSpan str pos end = takeSpan' end (str,"", pos)
takeSpan' :: RealSrcLoc -> (StringBuffer, String, RealSrcLoc) -> (StringBuffer, String, RealSrcLoc)
takeSpan' end (sb, taken, pos) | (srcLocLine pos `compare` srcLocLine end) `thenCmp` (srcLocCol pos `compare` srcLocCol end) == LT && not (atEnd sb)
= let (c,rem) = nextChar sb in takeSpan' end (rem, c:taken, advanceSrcLoc pos c)
takeSpan' _ (rem, taken, pos) = (rem, taken, pos)
thenCmp EQ o2 = o2
thenCmp o1 _ = o1
applyFragments :: SourceInfoTraversal node => [String] -> Ann node dom RngTemplateStage
-> Ann node dom SrcTemplateStage
applyFragments srcs = flip evalState srcs
. sourceInfoTraverseDown (SourceInfoTrf
(\ni -> do template <- mapM getTextFor (ni ^. rngTemplateNodeElems)
return $ SourceTemplateNode (RealSrcSpan $ ni ^. rngTemplateNodeRange) template 0 Nothing)
(\(RangeTemplateList rng bef aft sep indented seps)
-> do (own, rest) <- splitAt (length seps) <$> get
put rest
return (SourceTemplateList (RealSrcSpan rng) bef aft sep indented own 0 Nothing))
(\(RangeTemplateOpt rng bef aft) -> return (SourceTemplateOpt (RealSrcSpan rng) bef aft 0 Nothing)))
(return ()) (return ())
where getTextFor RangeChildElem = return ChildElem
getTextFor (RangeElem _) = do (src:rest) <- get
put rest
return (TextElem src)