content: {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE PatternGuards #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE ViewPatterns #-}
  {-# LANGUAGE RankNTypes #-}
  
  {- |
  Module      : SAWCore.ExternalFormat
  Copyright   : Galois, Inc. 2012-2015
  License     : BSD3
  Maintainer  : huffman@galois.com
  Stability   : experimental
  Portability : non-portable (language extensions)
  -}
  module SAWCore.ExternalFormat (
    scWriteExternal, scReadExternal
    ) where
  
  import qualified Control.Monad.State.Strict as State
  import Control.Monad.Trans.Class (MonadTrans(..))
  import Data.Map (Map)
  import qualified Data.Map as Map
  import qualified Data.Text as Text
  import Data.Text (Text)
  import qualified Data.Vector as V
  import Text.Read (readMaybe)
  import Text.URI
  
  import SAWCore.Name
  import SAWCore.SharedTerm
  import SAWCore.Term.Functor
  
  --------------------------------------------------------------------------------
  -- External text format
  
  type WriteM = State.State (Map TermIndex Int, Map VarIndex (Either Text NameInfo), [String], Int)
  
  renderNames :: Map VarIndex (Either Text NameInfo) -> String
  renderNames nms = show
    [ (idx, f nmi)
    | (idx,nmi) <- Map.toList nms
    ]
   where
     f (Left s) = Left s
     f (Right (ModuleIdentifier i))  = Right (Left (show i))
     f (Right (ImportedName uri as)) = Right (Right (render uri, as))
  
  readNames :: String -> Maybe (Map VarIndex (Either Text NameInfo))
  readNames xs = Map.fromList <$> (mapM readName =<< readMaybe xs)
   where
     readName :: (VarIndex, Either Text (Either Text (Text,[Text]))) -> Maybe (VarIndex, Either Text NameInfo)
     readName (idx, Left x) = pure (idx, Left x)
     readName (idx, Right (Left i)) = pure (idx, Right (ModuleIdentifier (parseIdent (Text.unpack i))))
     readName (idx, Right (Right (uri,as))) =
         do uri' <- mkURI uri
            pure (idx, Right (ImportedName uri' as))
  
  -- | Render to external text format
  scWriteExternal :: Term -> String
  scWriteExternal t0 =
      let (x, (_, nms, lns, _)) = State.runState (go t0) (Map.empty, Map.empty, [], 1)
      in unlines $
          [ unwords ["SAWCoreTerm", show x]
          , renderNames nms
          ] ++ reverse lns
    where
      nextId :: WriteM Int
      nextId =
         do (m, nms, lns, x) <- State.get
            State.put (m, nms, lns, x+1)
            return x
      output :: String -> WriteM ()
      output l =
         do (m, nms, lns, x) <- State.get
            State.put (m, nms, l:lns, x)
      memoize :: TermIndex -> WriteM Int
      memoize i =
         do (m, nms, lns, x) <- State.get
            State.put (Map.insert i x m, nms, lns, x+1)
            return x
      stashName :: Name -> WriteM ()
      stashName ec =
         do (m, nms, lns, x) <- State.get
            State.put (m, Map.insert (nameIndex ec) (Right (nameInfo ec)) nms, lns, x)
      stashVarName :: VarName -> WriteM ()
      stashVarName vn =
         do (m, nms, lns, x) <- State.get
            State.put (m, Map.insert (vnIndex vn) (Left (vnName vn)) nms, lns, x)
  
      go :: Term -> WriteM Int
      go (Unshared tf) = do
        tf' <- traverse go tf
        body <- writeTermF tf'
        x <- nextId
        output (unwords [show x, body])
        return x
  
      go STApp{ stAppIndex = i, stAppTermF = tf } = do
        (memo, _, _, _) <- State.get
        case Map.lookup i memo of
          Just x -> return x
          Nothing -> do
            tf' <- traverse go tf
            body <- writeTermF tf'
            x <- memoize i
            output (unwords [show x, body])
            return x
  
      writeTermF :: TermF Int -> WriteM String
      writeTermF tf =
        case tf of
          App e1 e2      -> pure $ unwords ["App", show e1, show e2]
          Lambda s t e   ->
            do stashVarName s
               pure $ unwords ["Lam", show (vnIndex s), show t, show e]
          Pi s t e       ->
            do stashVarName s
               pure $ unwords ["Pi", show (vnIndex s), show t, show e]
          Constant nm    ->
              do stashName nm
                 pure $ unwords ["Constant", show (nameIndex nm)]
          Variable nm tp ->
             do stashVarName nm
                pure $ unwords ["Variable", show (vnIndex nm), show tp]
          FTermF ftf     ->
            case ftf of
              UnitValue           -> pure $ unwords ["Unit"]
              UnitType            -> pure $ unwords ["UnitT"]
              PairValue x y       -> pure $ unwords ["Pair", show x, show y]
              PairType x y        -> pure $ unwords ["PairT", show x, show y]
              PairLeft e          -> pure $ unwords ["ProjL", show e]
              PairRight e         -> pure $ unwords ["ProjR", show e]
  
              Recursor (CompiledRecursor d s nps nixs ctorOrder) ->
                do stashName d
                   mapM_ stashName ctorOrder
                   let show_s = if s == propSort then "Prop" else drop 5 (show s)
                   pure $ unwords
                        (["Recursor", show (nameIndex d), show_s, show nps, show nixs
                         , show (map nameIndex ctorOrder)
                         ])
  
              RecordType elem_tps -> pure $ unwords ["RecordType", show elem_tps]
              RecordValue elems   -> pure $ unwords ["Record", show elems]
              RecordProj e prj    -> pure $ unwords ["RecordProj", show e, Text.unpack prj]
              Sort s h
                | s == propSort -> pure $ unwords ("Prop" : map show (sortFlagsToList h))
                | otherwise     -> pure $ unwords ("Sort" : drop 5 (show s) : map show (sortFlagsToList h))
                                                          -- /\ Ugly hack to drop "sort "
              NatLit n            -> pure $ unwords ["Nat", show n]
              ArrayValue e v      -> pure $ unwords ("Array" : show e :
                                              map show (V.toList v))
              StringLit s         -> pure $ unwords ["String", show s]
  
  
  -- | During parsing, we maintain two maps used for renumbering: The
  -- first is for the 'Int' values that appear in the external core
  -- file, and the second is for the 'VarIndex' values that appear
  -- inside 'Constant' and 'Variable' constructors. We do not reuse any
  -- such numbers that appear in the external file, but generate fresh
  -- ones that are valid in the current 'SharedContext'.
  type ReadM = State.StateT (Map Int Term, Map VarIndex (Either Text NameInfo), Map VarIndex VarIndex) IO
  
  scReadExternal :: SharedContext -> String -> IO Term
  scReadExternal sc input =
    case lines input of
      ( (words -> ["SAWCoreTerm", final]) : nmlist : rows ) ->
        case readNames nmlist of
          Nothing -> fail "scReadExternal: failed to parse name table"
          Just nms ->
            State.evalStateT (mapM_ (go . words) rows >> readIdx final) (Map.empty, nms, Map.empty)
  
      _ -> fail "scReadExternal: failed to parse input file"
    where
      go :: [String] -> ReadM ()
      go (tok : tokens) =
        do i <- readM tok
           tf <- parse tokens
           t <- lift $ scTermF sc tf
           (ts, nms, vs) <- State.get
           State.put (Map.insert i t ts, nms, vs)
      go [] = pure () -- empty lines are ignored
  
      readM :: forall a. Read a => String -> ReadM a
      readM tok =
        case readMaybe tok of
          Nothing -> fail $ "scReadExternal: parse error: " ++ show tok
          Just x -> pure x
  
      getTerm :: Int -> ReadM Term
      getTerm i =
        do (ts,_,_) <- State.get
           case Map.lookup i ts of
             Nothing -> fail $ "scReadExternal: invalid term index: " ++ show i
             Just t -> pure t
  
      readIdx :: String -> ReadM Term
      readIdx tok = getTerm =<< readM tok
  
      readCtorList :: String -> ReadM [Name]
      readCtorList str =
        do (ls :: [VarIndex]) <- readM str
           mapM readName' ls
  
      readName' :: VarIndex -> ReadM Name
      readName' vi =
        do (ts, nms, vs) <- State.get
           nmi <- case Map.lookup vi nms of
                    Just (Right nmi) -> pure nmi
                    _ -> lift $ fail $ "scReadExternal: ExtCns missing name info: " ++ show vi
           case nmi of
             ModuleIdentifier ident ->
               lift (scResolveNameByURI sc (moduleIdentToURI ident)) >>= \case
                 Just vi' -> pure (Name vi' nmi)
                 Nothing  -> lift $ fail $ "scReadExternal: missing module identifier: " ++ show ident
             ImportedName uri _aliases ->
               lift (scResolveNameByURI sc uri) >>= \case
                 Just vi' -> pure (Name vi' nmi)
                 Nothing -> case Map.lookup vi vs of
                   Just vi' -> pure $ Name vi' nmi
                   Nothing ->
                     do nm <- lift $ scRegisterName sc nmi
                        State.put (ts, nms, Map.insert vi (nameIndex nm) vs)
                        pure nm
  
      readName :: String -> ReadM Name
      readName i =
        do vi <- readM i
           readName' vi
  
      readVarName' :: VarIndex -> ReadM VarName
      readVarName' vi =
        do (ts, nms, vs) <- State.get
           case Map.lookup vi nms of
             Just (Left x) ->
               case Map.lookup vi vs of
                 Just vi' -> pure (VarName vi' x)
                 Nothing ->
                   do vn <- lift $ scFreshVarName sc x
                      State.put (ts, nms, Map.insert vi (vnIndex vn) vs)
                      pure vn
             _ -> lift $ fail $ "scReadExternal: VarName missing name: " ++ show vi
  
      readVarName :: String -> ReadM VarName
      readVarName i =
        do vi <- readM i
           readVarName' vi
  
      parse :: [String] -> ReadM (TermF Term)
      parse tokens =
        case tokens of
          ["App", e1, e2]     -> App <$> readIdx e1 <*> readIdx e2
          ["Lam", s, t, e]    -> Lambda <$> readVarName s <*> readIdx t <*> readIdx e
          ["Pi", s, t, e]     -> Pi <$> readVarName s <*> readIdx t <*> readIdx e
          ["Constant",i]      -> Constant <$> readName i
          ["ConstantOpaque",i]  -> Constant <$> readName i
          ["Unit"]            -> pure $ FTermF UnitValue
          ["UnitT"]           -> pure $ FTermF UnitType
          ["Pair", x, y]      -> FTermF <$> (PairValue <$> readIdx x <*> readIdx y)
          ["PairT", x, y]     -> FTermF <$> (PairType <$> readIdx x <*> readIdx y)
          ["ProjL", x]        -> FTermF <$> (PairLeft <$> readIdx x)
          ["ProjR", x]        -> FTermF <$> (PairRight <$> readIdx x)
  
          ["Recursor", i, s, nps, nixs, ctorOrder] ->
              do crec <- CompiledRecursor <$>
                          readName i <*>
                          (if s == "Prop" then pure propSort else mkSort <$> readM s) <*>
                          pure (read nps) <*>
                          pure (read nixs) <*>
                          readCtorList ctorOrder
                 pure (FTermF (Recursor crec))
  
          ["RecordType", elem_tps] ->
            FTermF <$> (RecordType <$> (traverse (traverse getTerm) =<< readM elem_tps))
          ["Record", elems] ->
            FTermF <$> (RecordValue <$> (traverse (traverse getTerm) =<< readM elems))
          ["RecordProj", e, prj] -> FTermF <$> (RecordProj <$> readIdx e <*> pure (Text.pack prj))
          ("Prop" : h)        -> FTermF <$> (Sort propSort . sortFlagsFromList <$> (mapM readM h))
          ("Sort" : s : h)    -> FTermF <$> (Sort <$> (mkSort <$> readM s) <*> (sortFlagsFromList <$> mapM readM h))
          ["Nat", n]          -> FTermF <$> (NatLit <$> readM n)
          ("Array" : e : es)  -> FTermF <$> (ArrayValue <$> readIdx e <*> (V.fromList <$> traverse readIdx es))
          ("String" : ts)     -> FTermF <$> (StringLit <$> (readM (unwords ts)))
          ["Variable", i, t]  -> Variable <$> readVarName i <*> readIdx t
          _ -> fail $ "Parse error: " ++ unwords tokens
  