<commit_before>{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE EmptyCase #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Verifier.SAW.Heapster.RustTypes where

import Prelude hiding (span)

import Data.Maybe
import Data.List hiding (span)
import GHC.TypeLits
import Data.Functor.Constant
import Data.Functor.Product
import Control.Applicative
import Control.Monad.Reader
import Control.Monad.Except
import Control.Monad.Trans.Maybe
import qualified Control.Monad.Fail as Fail

import Data.Parameterized.BoolRepr
import Data.Parameterized.Some
import Data.Parameterized.Context (Assignment, IsAppend(..),
                                   incSize, zeroSize, sizeInt,
                                   size, generate, extend)
import qualified Data.Parameterized.Context as Ctx

import Data.Binding.Hobbits
import Data.Binding.Hobbits.MonadBind
import qualified Data.Type.RList as RL
import qualified Data.Binding.Hobbits.NameSet as NameSet

import Language.Rust.Syntax
import Language.Rust.Parser
import qualified Language.Rust.Pretty as RustPP
import Language.Rust.Data.Ident (Ident(..), name)

import Prettyprinter as PP

import Lang.Crucible.Types
import Lang.Crucible.LLVM.Bytes
import Lang.Crucible.LLVM.MemModel hiding (Mutability(..))

import Verifier.SAW.Heapster.CruUtil
import Verifier.SAW.Heapster.Permissions


----------------------------------------------------------------------
-- * Helper Definitions for Translation
----------------------------------------------------------------------

-- | A permission of some llvm pointer type
data SomeLLVMPerm =
  forall w. (1 <= w, KnownNat w) => SomeLLVMPerm (ValuePerm (LLVMPointerType w))

$(mkNuMatching [t| SomeLLVMPerm |])

-- | Info used for converting Rust types to shapes
-- NOTE: @rciRecType@ should probably have some info about lifetimes
data RustConvInfo =
  RustConvInfo { rciPermEnv :: PermEnv,
                 rciCtx :: [(String, TypedName)],
                 rciRecType :: Maybe (RustName, [RustName], TypedName) }

-- | The default, top-level 'RustConvInfo' for a given 'PermEnv'
mkRustConvInfo :: PermEnv -> RustConvInfo
mkRustConvInfo env =
  RustConvInfo { rciPermEnv = env, rciCtx = [], rciRecType = Nothing }

-- | The Rust conversion monad is just a state-error monad
newtype RustConvM a =
  RustConvM { unRustConvM :: ReaderT RustConvInfo (Except String) a }
  deriving (Functor, Applicative, Monad,
            MonadError String, MonadReader RustConvInfo)

instance Fail.MonadFail RustConvM where
  fail = RustConvM . throwError

instance MonadBind RustConvM where
  mbM mb_m = RustConvM $ mbM $ fmap unRustConvM mb_m

-- | Prefix any error message with location information
atSpan :: Span -> RustConvM a -> RustConvM a
atSpan span m =
  catchError m (\msg -> fail ("At " ++ show span ++ ": " ++ msg))

-- | Run a Rust conversion computation with the given 'RustConvInfo', lifting
-- any errors into the outer monad
runLiftRustConvM :: Fail.MonadFail m => RustConvInfo -> RustConvM a -> m a
runLiftRustConvM info (RustConvM m) =
  case runExcept (runReaderT m info) of
    Left err -> fail err
    Right x -> return x

-- | Look up the 'TypedName' associated with a 'String' name
lookupTypedName :: String -> RustConvM TypedName
lookupTypedName str =
  fmap (lookup str) (rciCtx <$> ask) >>= \case
  Just n -> return n
  Nothing -> fail ("Could not find variable: " ++ show str)

-- | Look up a 'Name' with a given type
lookupName :: String -> TypeRepr a -> RustConvM (Name a)
lookupName str tp =
  lookupTypedName str >>= \n -> castTypedM "variable" tp n

-- | Build a 'PPInfo' structure for the names currently in scope
rsPPInfo :: RustConvM PPInfo
rsPPInfo =
  foldr (\(str, Some (Typed _ n)) -> ppInfoAddExprName str n) emptyPPInfo <$>
  rciCtx <$> ask

-- | The conversion of a context of Rust type and lifetime variables
type RustCtx = RAssign (Product (Constant String) TypeRepr)

-- | Build a 'RustCtx' for a single variable
rustCtx1 :: String -> TypeRepr a -> RustCtx (RNil :> a)
rustCtx1 name tp = MNil :>: Pair (Constant name) tp

-- | Extract a 'CruCtx' from a 'RustCtx'
rustCtxCtx :: RustCtx ctx -> CruCtx ctx
rustCtxCtx = cruCtxOfTypes . RL.map (\(Pair _ tp) -> tp)

-- | Extend a 'RustCtx' with a single binding on the right
rustCtxCons :: RustCtx ctx -> String -> TypeRepr a -> RustCtx (ctx :> a)
rustCtxCons ctx nm tp = ctx :>: Pair (Constant nm) tp

-- | Build a 'RustCtx' from the given variable names, all having the same type
rustCtxOfNames :: TypeRepr a -> [String] -> Some RustCtx
rustCtxOfNames tp =
  foldl (\(Some ctx) name -> Some (ctx :>: Pair (Constant name) tp)) (Some MNil)

-- | Run a 'RustConvM' computation in a context of bound type-level variables,
-- where the bound names are passed to the computation
inRustCtxF :: NuMatching a => RustCtx ctx -> (RAssign Name ctx -> RustConvM a) ->
              RustConvM (Mb ctx a)
inRustCtxF ctx m =
  mbM $ nuMulti (RL.map (\_ -> Proxy) ctx) $ \ns ->
  let ns_ctx =
        RL.toList $ RL.map2 (\n (Pair (Constant str) tp) ->
                              Constant (str, Some (Typed tp n))) ns ctx in
  local (\info -> info { rciCtx = ns_ctx ++ rciCtx info }) (m ns)

-- | Run a 'RustConvM' computation in a context of bound type-level variables
inRustCtx :: NuMatching a => RustCtx ctx -> RustConvM a ->
             RustConvM (Mb ctx a)
inRustCtx ctx m = inRustCtxF ctx (const m)

-- | Class for a generic "conversion from Rust" function, given the bit width of
-- the pointer type
class RsConvert w a b | w a -> b where
  rsConvert :: (1 <= w, KnownNat w) => prx w -> a -> RustConvM b


----------------------------------------------------------------------
-- * Converting Named Rust Types
----------------------------------------------------------------------

-- FIXME HERE: I think we no longer need SomeShapeFun

-- | A shape function with existentially quantified arguments
data SomeShapeFun w =
  forall args. SomeShapeFun (CruCtx args) (Mb args (PermExpr (LLVMShapeType w)))

-- | A sequence of 'PermExprs' along with their types
type TypedPermExprs = RAssign (Typed PermExpr)

-- | The empty sequence of typed permission expressions
emptyTypedPermExprs :: Some TypedPermExprs
emptyTypedPermExprs = Some MNil

appendTypedExprs :: Some TypedPermExprs -> Some TypedPermExprs ->
                    Some TypedPermExprs
appendTypedExprs (Some exprs1) (Some exprs2) = Some (RL.append exprs1 exprs2)

-- | Extract a type context from a 'TypedPermExprs'
typedPermExprsCtx :: TypedPermExprs ctx -> CruCtx ctx
typedPermExprsCtx = cruCtxOfTypes . RL.map typedType

-- | Extract the expressions from a 'TypedPermExprs'
typedPermExprsExprs :: TypedPermExprs ctx -> PermExprs ctx
typedPermExprsExprs = rassignToExprs . RL.map typedObj

-- | Convert a list of epxressions of a given type to a TypedPermExprs
typedExprsOfList :: TypeRepr a -> [PermExpr a] -> Some TypedPermExprs
typedExprsOfList tp =
  foldl (\(Some es) e -> Some (es :>: Typed tp e)) (Some MNil)

-- | Apply a 'SomeShapeFun', if possible
tryApplySomeShapeFun :: SomeShapeFun w -> Some TypedPermExprs ->
                        Maybe (PermExpr (LLVMShapeType w))
tryApplySomeShapeFun (SomeShapeFun ctx mb_sh) (Some exprs)
  | Just Refl <- testEquality ctx (typedPermExprsCtx exprs) =
    Just $ subst (substOfExprs $ typedPermExprsExprs exprs) mb_sh
tryApplySomeShapeFun _ _ = Nothing

-- | Build a 'SomeShapeFun' with no arguments
constShapeFun :: PermExpr (LLVMShapeType w) -> SomeShapeFun w
constShapeFun sh = SomeShapeFun CruCtxNil (emptyMb sh)

-- | Build the shape @fieldsh(exists z:bv sz.eq(llvmword(z))@
sizedIntShapeFun :: (1 <= w, KnownNat w, 1 <= sz, KnownNat sz) =>
                    prx1 w -> prx2 sz -> SomeShapeFun w
sizedIntShapeFun _ sz =
  constShapeFun $ PExpr_FieldShape (LLVMFieldShape $
                                    ValPerm_Exists $ llvmExEqWord sz)

-- | Build a `SomeShapeFun` from `SomeNamedShape`
namedShapeShapeFun :: NatRepr w -> SomeNamedShape -> RustConvM (SomeShapeFun w)
namedShapeShapeFun w (SomeNamedShape nmsh)
  | Just Refl <- testEquality w (natRepr nmsh) = return $
    SomeShapeFun (namedShapeArgs nmsh)
                 (nuMulti (cruCtxProxies (namedShapeArgs nmsh))
                          (\ns -> PExpr_NamedShape Nothing Nothing nmsh (namesToExprs ns)))
namedShapeShapeFun w (SomeNamedShape nmsh) =
  fail $ renderDoc $ fillSep
  [pretty "Incorrect size of shape" <+> pretty (namedShapeName nmsh),
   pretty "Expected:" <+> pretty (intValue w),
   pretty "Actual:" <+> pretty (intValue (natRepr nmsh))]

-- | A fully qualified Rust path without any of the parameters; e.g.,
-- @Foo<X>::Bar<Y,Z>::Baz@ just becomes @[Foo,Bar,Baz]@
newtype RustName = RustName [Ident] deriving (Eq)

-- | Convert a 'RustName' to a string by interspersing "::"
flattenRustName :: RustName -> String
flattenRustName (RustName ids) = concat $ intersperse "::" $ map name ids

instance Show RustName where
  show = show . flattenRustName

instance RsConvert w RustName (SomeShapeFun w) where
  rsConvert w nm =
    do let str = flattenRustName nm
       env <- rciPermEnv <$> ask
       case lookupNamedShape env str of
         Just nmsh -> namedShapeShapeFun (natRepr w) nmsh
         Nothing ->
           do n <- lookupName str (LLVMShapeRepr (natRepr w))
              return $ constShapeFun (PExpr_Var n)

-- | Get the "name" = sequence of identifiers out of a Rust path
rsPathName :: Path a -> RustName
rsPathName (Path _ segments _) =
  RustName $ map (\(PathSegment rust_id _ _) -> rust_id) segments

-- | Get all the parameters out of a Rust path
rsPathParams :: Path a -> [PathParameters a]
rsPathParams (Path _ segments _) =
  mapMaybe (\(PathSegment _ maybe_params _) -> maybe_params) segments

-- | Get the 'RustName' of a type, if it's a 'PathTy'
tyName :: Ty a -> Maybe RustName
tyName (PathTy _ path _) = Just $ rsPathName path
tyName _ = Nothing

-- | Decide whether a Rust type is named (i.e. a 'PathTy')
isNamedType :: Ty a -> Bool
isNamedType (PathTy _ _ _) = True
isNamedType _ = False

-- | Decide whether 'PathParameters' are all named types (angle-bracketed only)
isNamedParams :: PathParameters a -> Bool
isNamedParams (AngleBracketed _ tys _ _) = all isNamedType tys
isNamedParams _ = error "Parenthesized types not supported"

-- | Get all of the 'RustName's of path parameters, if they're angle-bracketed
pParamNames :: PathParameters a -> [RustName]
pParamNames (AngleBracketed _ tys _ _) = mapMaybe tyName tys
pParamNames _ = error "Parenthesized types not supported"

-- | Modify a 'RustConvM' to be run with a recursive type
withRecType :: (1 <= w, KnownNat w) => RustName -> [RustName] -> Name (LLVMShapeType w) ->
               RustConvM a -> RustConvM a
withRecType rust_n rust_ns rec_n = local (\info -> info { rciRecType = Just (rust_n, rust_ns, Some (Typed knownRepr rec_n)) })


----------------------------------------------------------------------
-- * Converting Rust Types to Heapster Shapes
----------------------------------------------------------------------

-- | Test if a shape matches the translation of a slice type, and, if so, return
-- the stride and the fields of the slice, where the latter can have the length
-- free
matchSliceShape :: PermExpr (LLVMShapeType w) ->
                   Maybe (Bytes,
                          Binding (BVType w) (PermExpr (LLVMShapeType w)))
matchSliceShape (PExpr_ExShape
                 [nuP| PExpr_ArrayShape (PExpr_Var len) stride mb_sh |])
  | Left Member_Base <- mbNameBoundP len =
    Just (mbLift stride, mb_sh)
matchSliceShape (PExpr_NamedShape _ _ nmsh@(NamedShape _ _
                                            (DefinedShapeBody _)) args) =
  matchSliceShape (unfoldNamedShape nmsh args)
matchSliceShape _ = Nothing

-- Convert a 'Mutability' to a modality override for a 'PExpr_PtrShape'; mutable
-- references inherit the modality of the container they are in, so they
-- translate to 'Nothing'
instance RsConvert w Mutability (Maybe (PermExpr RWModalityType)) where
  rsConvert _ Mutable = return Nothing
  rsConvert _ Immutable = return (Just PExpr_Read)

instance RsConvert w (Lifetime Span) (PermExpr LifetimeType) where
  rsConvert _ (Lifetime "static" _) = return PExpr_Always
  rsConvert _ (Lifetime l span) =
    PExpr_Var <$> atSpan span (lookupName l knownRepr)

instance RsConvert w (PathParameters Span) (Some TypedPermExprs) where
  rsConvert w (AngleBracketed rust_ls rust_tps [] _) =
    do ls <- mapM (rsConvert w) rust_ls
       shs <- mapM (rsConvert w) rust_tps
       return $ appendTypedExprs
         (typedExprsOfList knownRepr ls) (typedExprsOfList knownRepr shs)
  rsConvert _ (AngleBracketed _ _ (_:_) _) =
    error "rsConvert: angle-bracketed arguments not supported"
  rsConvert _ (Parenthesized _ _ _) =
    error "rsConvert: parenthesized types not supported"

instance RsConvert w [PathParameters Span] (Some TypedPermExprs) where
  rsConvert w paramss =
    foldr appendTypedExprs emptyTypedPermExprs <$> mapM (rsConvert w) paramss

instance RsConvert w (Ty Span) (PermExpr (LLVMShapeType w)) where
  rsConvert w (Slice tp _) =
    do sh <- rsConvert w tp
       case llvmShapeLength sh of
         Just (bvMatchConstInt -> Just stride) ->
           return (PExpr_ExShape $ nu $ \n ->
                    PExpr_ArrayShape (PExpr_Var n) (fromInteger stride) sh)
         _ ->
           rsPPInfo >>= \ppInfo ->
           fail ("rsConvert: slices not supported for dynamically-sized type: "
                 ++ show (RustPP.pretty tp) ++ " with translation:\n"
                 ++ renderDoc (permPretty ppInfo sh))
  rsConvert _ (Rptr Nothing _ _ _) =
    fail "rsConvert: lifetimes must be supplied for reference types"
  rsConvert w (Rptr (Just rust_l) mut tp' _) =
    do l <- rsConvert w rust_l
       sh <- rsConvert w tp'
       rw <- rsConvert w mut
       case sh of
         -- Test if sh is a slice type = an array of existential length
         (matchSliceShape -> Just (stride,mb_sh)) ->
           -- If so, build a "fat pointer" = a pair of a pointer to our array
           -- shape plus a length value
           return $ PExpr_ExShape $ nu $ \n ->
           PExpr_SeqShape (PExpr_PtrShape rw (Just l) $
                           PExpr_ArrayShape (PExpr_Var n) stride $
                           subst1 (PExpr_Var n) mb_sh)
           (PExpr_FieldShape $ LLVMFieldShape $ ValPerm_Eq $
            PExpr_LLVMWord $ PExpr_Var n)

         -- If it's not a slice, make sure it has a known size
         _ | Just len <- llvmShapeLength sh
           , isJust (bvMatchConst len) ->
             return $ PExpr_PtrShape rw (Just l) sh

         -- Otherwise, it's a non-standard dynamically-sized type, which we
         -- don't quite know how to handle yet...
         _ -> fail "rsConvert: pointer to non-slice dynamically-sized type"
  rsConvert w (PathTy Nothing path _) =
    do mrec <- asks rciRecType
       case mrec of
         Just (rec_n, rec_arg_ns, sh_nm)
           | rec_n == rsPathName path &&
             all isNamedParams (rsPathParams path) &&
             rec_arg_ns == concatMap pParamNames (rsPathParams path) ->
             PExpr_Var <$> castTypedM "TypedName" (LLVMShapeRepr (natRepr w)) sh_nm
         Just (rec_n, _, _)
           | rec_n == rsPathName path -> fail "Arguments do not match"
         _ ->
           do someShapeFn@(SomeShapeFun expected _ ) <- rsConvert w (rsPathName path)
              someTypedArgs@(Some tyArgs) <- rsConvert w (rsPathParams path)
              let actual = typedPermExprsCtx tyArgs
              case tryApplySomeShapeFun someShapeFn someTypedArgs of
                Just shTp -> return shTp
                Nothing ->
                  fail $ renderDoc $ fillSep
                  [ pretty "Converting PathTy: " <+> pretty (show $ rsPathName path)
                  , pretty "Expected arguments:" <+> pretty expected
                  , pretty "Actual arguments:" <+> pretty actual
                  ]
  rsConvert (w :: prx w) (BareFn _ abi rust_ls2 fn_tp span) =
    do Some3FunPerm fun_perm <- rsConvertMonoFun w span abi rust_ls2 fn_tp
       let args = funPermArgs fun_perm
       case cruCtxToReprEq args of
         Refl ->
           return $ PExpr_FieldShape $ LLVMFieldShape @w @w $ ValPerm_Conj1 $
           Perm_LLVMFunPtr
           (FunctionHandleRepr (cruCtxToRepr args) (funPermRet fun_perm)) $
           ValPerm_Conj1 $ Perm_Fun fun_perm
  rsConvert w (TupTy tys _) =
    do tyShs <- mapM (rsConvert w) tys
       return $ foldr PExpr_SeqShape PExpr_EmptyShape tyShs
  rsConvert _ (Never _) =
    return $ PExpr_FieldShape $ LLVMFieldShape @w @w ValPerm_False
  rsConvert _ tp = fail ("Rust type not supported: " ++ show tp)

instance RsConvert w (Arg Span) (PermExpr (LLVMShapeType w)) where
  rsConvert w (Arg _ tp _) = rsConvert w tp
  rsConvert _ _ = error "rsConvert (Arg): argument form not yet handled"

instance RsConvert w (Generics Span) (Some RustCtx) where
  rsConvert w (Generics ltdefs tyvars _ _) =
    return $ foldl addTyVar (foldl addLt (Some MNil) ltdefs) tyvars
    where
      addLt (Some ctx) ltdef =
        Some (ctx :>: Pair (Constant (lifetimeDefName ltdef)) LifetimeRepr)

      addTyVar (Some ctx) tyvar =
        Some (ctx :>: Pair (Constant (tyParamName tyvar)) (LLVMShapeRepr (natRepr w)))

-- | Return true if and only if the provided Rust type definition is recursive
isRecursiveDef :: Item Span -> Bool
isRecursiveDef item =
  case item of
    Enum _ _ n variants _ _ -> any (containsName n . getVD) variants
    StructItem _ _ n vd _ _ -> containsName n vd
    _ -> False<commit_msg>Use the shiny new false shape for the Rust ! type<commit_after>{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE EmptyCase #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Verifier.SAW.Heapster.RustTypes where

import Prelude hiding (span)

import Data.Maybe
import Data.List hiding (span)
import GHC.TypeLits
import Data.Functor.Constant
import Data.Functor.Product
import Control.Applicative
import Control.Monad.Reader
import Control.Monad.Except
import Control.Monad.Trans.Maybe
import qualified Control.Monad.Fail as Fail

import Data.Parameterized.BoolRepr
import Data.Parameterized.Some
import Data.Parameterized.Context (Assignment, IsAppend(..),
                                   incSize, zeroSize, sizeInt,
                                   size, generate, extend)
import qualified Data.Parameterized.Context as Ctx

import Data.Binding.Hobbits
import Data.Binding.Hobbits.MonadBind
import qualified Data.Type.RList as RL
import qualified Data.Binding.Hobbits.NameSet as NameSet

import Language.Rust.Syntax
import Language.Rust.Parser
import qualified Language.Rust.Pretty as RustPP
import Language.Rust.Data.Ident (Ident(..), name)

import Prettyprinter as PP

import Lang.Crucible.Types
import Lang.Crucible.LLVM.Bytes
import Lang.Crucible.LLVM.MemModel hiding (Mutability(..))

import Verifier.SAW.Heapster.CruUtil
import Verifier.SAW.Heapster.Permissions


----------------------------------------------------------------------
-- * Helper Definitions for Translation
----------------------------------------------------------------------

-- | A permission of some llvm pointer type
data SomeLLVMPerm =
  forall w. (1 <= w, KnownNat w) => SomeLLVMPerm (ValuePerm (LLVMPointerType w))

$(mkNuMatching [t| SomeLLVMPerm |])

-- | Info used for converting Rust types to shapes
-- NOTE: @rciRecType@ should probably have some info about lifetimes
data RustConvInfo =
  RustConvInfo { rciPermEnv :: PermEnv,
                 rciCtx :: [(String, TypedName)],
                 rciRecType :: Maybe (RustName, [RustName], TypedName) }

-- | The default, top-level 'RustConvInfo' for a given 'PermEnv'
mkRustConvInfo :: PermEnv -> RustConvInfo
mkRustConvInfo env =
  RustConvInfo { rciPermEnv = env, rciCtx = [], rciRecType = Nothing }

-- | The Rust conversion monad is just a state-error monad
newtype RustConvM a =
  RustConvM { unRustConvM :: ReaderT RustConvInfo (Except String) a }
  deriving (Functor, Applicative, Monad,
            MonadError String, MonadReader RustConvInfo)

instance Fail.MonadFail RustConvM where
  fail = RustConvM . throwError

instance MonadBind RustConvM where
  mbM mb_m = RustConvM $ mbM $ fmap unRustConvM mb_m

-- | Prefix any error message with location information
atSpan :: Span -> RustConvM a -> RustConvM a
atSpan span m =
  catchError m (\msg -> fail ("At " ++ show span ++ ": " ++ msg))

-- | Run a Rust conversion computation with the given 'RustConvInfo', lifting
-- any errors into the outer monad
runLiftRustConvM :: Fail.MonadFail m => RustConvInfo -> RustConvM a -> m a
runLiftRustConvM info (RustConvM m) =
  case runExcept (runReaderT m info) of
    Left err -> fail err
    Right x -> return x

-- | Look up the 'TypedName' associated with a 'String' name
lookupTypedName :: String -> RustConvM TypedName
lookupTypedName str =
  fmap (lookup str) (rciCtx <$> ask) >>= \case
  Just n -> return n
  Nothing -> fail ("Could not find variable: " ++ show str)

-- | Look up a 'Name' with a given type
lookupName :: String -> TypeRepr a -> RustConvM (Name a)
lookupName str tp =
  lookupTypedName str >>= \n -> castTypedM "variable" tp n

-- | Build a 'PPInfo' structure for the names currently in scope
rsPPInfo :: RustConvM PPInfo
rsPPInfo =
  foldr (\(str, Some (Typed _ n)) -> ppInfoAddExprName str n) emptyPPInfo <$>
  rciCtx <$> ask

-- | The conversion of a context of Rust type and lifetime variables
type RustCtx = RAssign (Product (Constant String) TypeRepr)

-- | Build a 'RustCtx' for a single variable
rustCtx1 :: String -> TypeRepr a -> RustCtx (RNil :> a)
rustCtx1 name tp = MNil :>: Pair (Constant name) tp

-- | Extract a 'CruCtx' from a 'RustCtx'
rustCtxCtx :: RustCtx ctx -> CruCtx ctx
rustCtxCtx = cruCtxOfTypes . RL.map (\(Pair _ tp) -> tp)

-- | Extend a 'RustCtx' with a single binding on the right
rustCtxCons :: RustCtx ctx -> String -> TypeRepr a -> RustCtx (ctx :> a)
rustCtxCons ctx nm tp = ctx :>: Pair (Constant nm) tp

-- | Build a 'RustCtx' from the given variable names, all having the same type
rustCtxOfNames :: TypeRepr a -> [String] -> Some RustCtx
rustCtxOfNames tp =
  foldl (\(Some ctx) name -> Some (ctx :>: Pair (Constant name) tp)) (Some MNil)

-- | Run a 'RustConvM' computation in a context of bound type-level variables,
-- where the bound names are passed to the computation
inRustCtxF :: NuMatching a => RustCtx ctx -> (RAssign Name ctx -> RustConvM a) ->
              RustConvM (Mb ctx a)
inRustCtxF ctx m =
  mbM $ nuMulti (RL.map (\_ -> Proxy) ctx) $ \ns ->
  let ns_ctx =
        RL.toList $ RL.map2 (\n (Pair (Constant str) tp) ->
                              Constant (str, Some (Typed tp n))) ns ctx in
  local (\info -> info { rciCtx = ns_ctx ++ rciCtx info }) (m ns)

-- | Run a 'RustConvM' computation in a context of bound type-level variables
inRustCtx :: NuMatching a => RustCtx ctx -> RustConvM a ->
             RustConvM (Mb ctx a)
inRustCtx ctx m = inRustCtxF ctx (const m)

-- | Class for a generic "conversion from Rust" function, given the bit width of
-- the pointer type
class RsConvert w a b | w a -> b where
  rsConvert :: (1 <= w, KnownNat w) => prx w -> a -> RustConvM b


----------------------------------------------------------------------
-- * Converting Named Rust Types
----------------------------------------------------------------------

-- FIXME HERE: I think we no longer need SomeShapeFun

-- | A shape function with existentially quantified arguments
data SomeShapeFun w =
  forall args. SomeShapeFun (CruCtx args) (Mb args (PermExpr (LLVMShapeType w)))

-- | A sequence of 'PermExprs' along with their types
type TypedPermExprs = RAssign (Typed PermExpr)

-- | The empty sequence of typed permission expressions
emptyTypedPermExprs :: Some TypedPermExprs
emptyTypedPermExprs = Some MNil

appendTypedExprs :: Some TypedPermExprs -> Some TypedPermExprs ->
                    Some TypedPermExprs
appendTypedExprs (Some exprs1) (Some exprs2) = Some (RL.append exprs1 exprs2)

-- | Extract a type context from a 'TypedPermExprs'
typedPermExprsCtx :: TypedPermExprs ctx -> CruCtx ctx
typedPermExprsCtx = cruCtxOfTypes . RL.map typedType

-- | Extract the expressions from a 'TypedPermExprs'
typedPermExprsExprs :: TypedPermExprs ctx -> PermExprs ctx
typedPermExprsExprs = rassignToExprs . RL.map typedObj

-- | Convert a list of epxressions of a given type to a TypedPermExprs
typedExprsOfList :: TypeRepr a -> [PermExpr a] -> Some TypedPermExprs
typedExprsOfList tp =
  foldl (\(Some es) e -> Some (es :>: Typed tp e)) (Some MNil)

-- | Apply a 'SomeShapeFun', if possible
tryApplySomeShapeFun :: SomeShapeFun w -> Some TypedPermExprs ->
                        Maybe (PermExpr (LLVMShapeType w))
tryApplySomeShapeFun (SomeShapeFun ctx mb_sh) (Some exprs)
  | Just Refl <- testEquality ctx (typedPermExprsCtx exprs) =
    Just $ subst (substOfExprs $ typedPermExprsExprs exprs) mb_sh
tryApplySomeShapeFun _ _ = Nothing

-- | Build a 'SomeShapeFun' with no arguments
constShapeFun :: PermExpr (LLVMShapeType w) -> SomeShapeFun w
constShapeFun sh = SomeShapeFun CruCtxNil (emptyMb sh)

-- | Build the shape @fieldsh(exists z:bv sz.eq(llvmword(z))@
sizedIntShapeFun :: (1 <= w, KnownNat w, 1 <= sz, KnownNat sz) =>
                    prx1 w -> prx2 sz -> SomeShapeFun w
sizedIntShapeFun _ sz =
  constShapeFun $ PExpr_FieldShape (LLVMFieldShape $
                                    ValPerm_Exists $ llvmExEqWord sz)

-- | Build a `SomeShapeFun` from `SomeNamedShape`
namedShapeShapeFun :: NatRepr w -> SomeNamedShape -> RustConvM (SomeShapeFun w)
namedShapeShapeFun w (SomeNamedShape nmsh)
  | Just Refl <- testEquality w (natRepr nmsh) = return $
    SomeShapeFun (namedShapeArgs nmsh)
                 (nuMulti (cruCtxProxies (namedShapeArgs nmsh))
                          (\ns -> PExpr_NamedShape Nothing Nothing nmsh (namesToExprs ns)))
namedShapeShapeFun w (SomeNamedShape nmsh) =
  fail $ renderDoc $ fillSep
  [pretty "Incorrect size of shape" <+> pretty (namedShapeName nmsh),
   pretty "Expected:" <+> pretty (intValue w),
   pretty "Actual:" <+> pretty (intValue (natRepr nmsh))]

-- | A fully qualified Rust path without any of the parameters; e.g.,
-- @Foo<X>::Bar<Y,Z>::Baz@ just becomes @[Foo,Bar,Baz]@
newtype RustName = RustName [Ident] deriving (Eq)

-- | Convert a 'RustName' to a string by interspersing "::"
flattenRustName :: RustName -> String
flattenRustName (RustName ids) = concat $ intersperse "::" $ map name ids

instance Show RustName where
  show = show . flattenRustName

instance RsConvert w RustName (SomeShapeFun w) where
  rsConvert w nm =
    do let str = flattenRustName nm
       env <- rciPermEnv <$> ask
       case lookupNamedShape env str of
         Just nmsh -> namedShapeShapeFun (natRepr w) nmsh
         Nothing ->
           do n <- lookupName str (LLVMShapeRepr (natRepr w))
              return $ constShapeFun (PExpr_Var n)

-- | Get the "name" = sequence of identifiers out of a Rust path
rsPathName :: Path a -> RustName
rsPathName (Path _ segments _) =
  RustName $ map (\(PathSegment rust_id _ _) -> rust_id) segments

-- | Get all the parameters out of a Rust path
rsPathParams :: Path a -> [PathParameters a]
rsPathParams (Path _ segments _) =
  mapMaybe (\(PathSegment _ maybe_params _) -> maybe_params) segments

-- | Get the 'RustName' of a type, if it's a 'PathTy'
tyName :: Ty a -> Maybe RustName
tyName (PathTy _ path _) = Just $ rsPathName path
tyName _ = Nothing

-- | Decide whether a Rust type is named (i.e. a 'PathTy')
isNamedType :: Ty a -> Bool
isNamedType (PathTy _ _ _) = True
isNamedType _ = False

-- | Decide whether 'PathParameters' are all named types (angle-bracketed only)
isNamedParams :: PathParameters a -> Bool
isNamedParams (AngleBracketed _ tys _ _) = all isNamedType tys
isNamedParams _ = error "Parenthesized types not supported"

-- | Get all of the 'RustName's of path parameters, if they're angle-bracketed
pParamNames :: PathParameters a -> [RustName]
pParamNames (AngleBracketed _ tys _ _) = mapMaybe tyName tys
pParamNames _ = error "Parenthesized types not supported"

-- | Modify a 'RustConvM' to be run with a recursive type
withRecType :: (1 <= w, KnownNat w) => RustName -> [RustName] -> Name (LLVMShapeType w) ->
               RustConvM a -> RustConvM a
withRecType rust_n rust_ns rec_n = local (\info -> info { rciRecType = Just (rust_n, rust_ns, Some (Typed knownRepr rec_n)) })


----------------------------------------------------------------------
-- * Converting Rust Types to Heapster Shapes
----------------------------------------------------------------------

-- | Test if a shape matches the translation of a slice type, and, if so, return
-- the stride and the fields of the slice, where the latter can have the length
-- free
matchSliceShape :: PermExpr (LLVMShapeType w) ->
                   Maybe (Bytes,
                          Binding (BVType w) (PermExpr (LLVMShapeType w)))
matchSliceShape (PExpr_ExShape
                 [nuP| PExpr_ArrayShape (PExpr_Var len) stride mb_sh |])
  | Left Member_Base <- mbNameBoundP len =
    Just (mbLift stride, mb_sh)
matchSliceShape (PExpr_NamedShape _ _ nmsh@(NamedShape _ _
                                            (DefinedShapeBody _)) args) =
  matchSliceShape (unfoldNamedShape nmsh args)
matchSliceShape _ = Nothing

-- Convert a 'Mutability' to a modality override for a 'PExpr_PtrShape'; mutable
-- references inherit the modality of the container they are in, so they
-- translate to 'Nothing'
instance RsConvert w Mutability (Maybe (PermExpr RWModalityType)) where
  rsConvert _ Mutable = return Nothing
  rsConvert _ Immutable = return (Just PExpr_Read)

instance RsConvert w (Lifetime Span) (PermExpr LifetimeType) where
  rsConvert _ (Lifetime "static" _) = return PExpr_Always
  rsConvert _ (Lifetime l span) =
    PExpr_Var <$> atSpan span (lookupName l knownRepr)

instance RsConvert w (PathParameters Span) (Some TypedPermExprs) where
  rsConvert w (AngleBracketed rust_ls rust_tps [] _) =
    do ls <- mapM (rsConvert w) rust_ls
       shs <- mapM (rsConvert w) rust_tps
       return $ appendTypedExprs
         (typedExprsOfList knownRepr ls) (typedExprsOfList knownRepr shs)
  rsConvert _ (AngleBracketed _ _ (_:_) _) =
    error "rsConvert: angle-bracketed arguments not supported"
  rsConvert _ (Parenthesized _ _ _) =
    error "rsConvert: parenthesized types not supported"

instance RsConvert w [PathParameters Span] (Some TypedPermExprs) where
  rsConvert w paramss =
    foldr appendTypedExprs emptyTypedPermExprs <$> mapM (rsConvert w) paramss

instance RsConvert w (Ty Span) (PermExpr (LLVMShapeType w)) where
  rsConvert w (Slice tp _) =
    do sh <- rsConvert w tp
       case llvmShapeLength sh of
         Just (bvMatchConstInt -> Just stride) ->
           return (PExpr_ExShape $ nu $ \n ->
                    PExpr_ArrayShape (PExpr_Var n) (fromInteger stride) sh)
         _ ->
           rsPPInfo >>= \ppInfo ->
           fail ("rsConvert: slices not supported for dynamically-sized type: "
                 ++ show (RustPP.pretty tp) ++ " with translation:\n"
                 ++ renderDoc (permPretty ppInfo sh))
  rsConvert _ (Rptr Nothing _ _ _) =
    fail "rsConvert: lifetimes must be supplied for reference types"
  rsConvert w (Rptr (Just rust_l) mut tp' _) =
    do l <- rsConvert w rust_l
       sh <- rsConvert w tp'
       rw <- rsConvert w mut
       case sh of
         -- Test if sh is a slice type = an array of existential length
         (matchSliceShape -> Just (stride,mb_sh)) ->
           -- If so, build a "fat pointer" = a pair of a pointer to our array
           -- shape plus a length value
           return $ PExpr_ExShape $ nu $ \n ->
           PExpr_SeqShape (PExpr_PtrShape rw (Just l) $
                           PExpr_ArrayShape (PExpr_Var n) stride $
                           subst1 (PExpr_Var n) mb_sh)
           (PExpr_FieldShape $ LLVMFieldShape $ ValPerm_Eq $
            PExpr_LLVMWord $ PExpr_Var n)

         -- If it's not a slice, make sure it has a known size
         _ | Just len <- llvmShapeLength sh
           , isJust (bvMatchConst len) ->
             return $ PExpr_PtrShape rw (Just l) sh

         -- Otherwise, it's a non-standard dynamically-sized type, which we
         -- don't quite know how to handle yet...
         _ -> fail "rsConvert: pointer to non-slice dynamically-sized type"
  rsConvert w (PathTy Nothing path _) =
    do mrec <- asks rciRecType
       case mrec of
         Just (rec_n, rec_arg_ns, sh_nm)
           | rec_n == rsPathName path &&
             all isNamedParams (rsPathParams path) &&
             rec_arg_ns == concatMap pParamNames (rsPathParams path) ->
             PExpr_Var <$> castTypedM "TypedName" (LLVMShapeRepr (natRepr w)) sh_nm
         Just (rec_n, _, _)
           | rec_n == rsPathName path -> fail "Arguments do not match"
         _ ->
           do someShapeFn@(SomeShapeFun expected _ ) <- rsConvert w (rsPathName path)
              someTypedArgs@(Some tyArgs) <- rsConvert w (rsPathParams path)
              let actual = typedPermExprsCtx tyArgs
              case tryApplySomeShapeFun someShapeFn someTypedArgs of
                Just shTp -> return shTp
                Nothing ->
                  fail $ renderDoc $ fillSep
                  [ pretty "Converting PathTy: " <+> pretty (show $ rsPathName path)
                  , pretty "Expected arguments:" <+> pretty expected
                  , pretty "Actual arguments:" <+> pretty actual
                  ]
  rsConvert (w :: prx w) (BareFn _ abi rust_ls2 fn_tp span) =
    do Some3FunPerm fun_perm <- rsConvertMonoFun w span abi rust_ls2 fn_tp
       let args = funPermArgs fun_perm
       case cruCtxToReprEq args of
         Refl ->
           return $ PExpr_FieldShape $ LLVMFieldShape @w @w $ ValPerm_Conj1 $
           Perm_LLVMFunPtr
           (FunctionHandleRepr (cruCtxToRepr args) (funPermRet fun_perm)) $
           ValPerm_Conj1 $ Perm_Fun fun_perm
  rsConvert w (TupTy tys _) =
    do tyShs <- mapM (rsConvert w) tys
       return $ foldr PExpr_SeqShape PExpr_EmptyShape tyShs
  rsConvert _ (Never _) =
    return $ PExpr_FalseShape
  rsConvert _ tp = fail ("Rust type not supported: " ++ show tp)

instance RsConvert w (Arg Span) (PermExpr (LLVMShapeType w)) where
  rsConvert w (Arg _ tp _) = rsConvert w tp
  rsConvert _ _ = error "rsConvert (Arg): argument form not yet handled"

instance RsConvert w (Generics Span) (Some RustCtx) where
  rsConvert w (Generics ltdefs tyvars _ _) =
    return $ foldl addTyVar (foldl addLt (Some MNil) ltdefs) tyvars
    where
      addLt (Some ctx) ltdef =
        Some (ctx :>: Pair (Constant (lifetimeDefName ltdef)) LifetimeRepr)

      addTyVar (Some ctx) tyvar =
        Some (ctx :>: Pair (Constant (tyParamName tyvar)) (LLVMShapeRepr (natRepr w)))

-- | Return true if and only if the provided Rust type definition is recursive
isRecursiveDef :: Item Span -> Bool
isRecursiveDef item =
  case item of
    Enum _ _ n variants _ _ -> any (containsName n . getVD) variants
    StructItem _ _ n vd _ _ -> containsName n vd
    _ -> False
{'fraction_non_alphanumeric': 0.11557307492266841, 'fraction_numerical': 0.0010949604445539404, 'mean_word_length': 2.782563677780079, 'pattern_counts': {'":': 6, '<': 77, '<?xml version=': 0, '>': 253, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5101937', 'n_tokens_mistral': 11593, 'n_tokens_neox': 11395, 'n_words': 3792}