diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,5 +1,10 @@ # Changelog for inline-asm +## v0.3.1.0 + +* Support passing `ByteString`s to the assembly. +* Introduced a helper synonym `RET_HASK` to get back into the Haskell land explicitly. + ## v0.3.0.0 * Support named arguments for doubles and floats as well. diff --git a/inline-asm.cabal b/inline-asm.cabal --- a/inline-asm.cabal +++ b/inline-asm.cabal @@ -4,10 +4,10 @@ -- -- see: https://github.com/sol/hpack -- --- hash: 07f84f8e6bb0bdb11a97883ebf8eb73d17aeafe6f29f05996646e7487b23583f +-- hash: 7cf75cf1921021ab02ccd85a363e971578c52fb35a1ce1a35575ea001544b62c name: inline-asm -version: 0.2.1.0 +version: 0.3.1.0 synopsis: Inline some Assembly in ur Haskell! description: Please see the README on GitHub at category: FFI @@ -40,6 +40,7 @@ ghc-options: -Wall build-depends: base >=4.7 && <5 + , bytestring , containers , either , ghc-prim @@ -58,6 +59,7 @@ ghc-options: -Wall -threaded -rtsopts -with-rtsopts=-N build-depends: base >=4.7 && <5 + , bytestring , containers , either , ghc-prim @@ -79,6 +81,7 @@ build-depends: QuickCheck , base >=4.7 && <5 + , bytestring , containers , either , ghc-prim diff --git a/src/Language/Asm/Inline.hs b/src/Language/Asm/Inline.hs --- a/src/Language/Asm/Inline.hs +++ b/src/Language/Asm/Inline.hs @@ -5,14 +5,20 @@ module Language.Asm.Inline(defineAsmFun) where +import qualified Data.ByteString as BS +import qualified Data.ByteString.Internal as BS import Control.Monad import Data.Generics.Uniplate.Data +import Data.List import Foreign.Ptr +import Foreign.ForeignPtr.Unsafe import GHC.Prim import GHC.Ptr import GHC.Types hiding (Type) +import GHC.Word import Language.Haskell.TH import Language.Haskell.TH.Syntax +import System.IO.Unsafe import Language.Asm.Inline.AsmCode import Language.Asm.Inline.Util @@ -41,12 +47,19 @@ unbox (Ptr p) = p rebox = Ptr +replace :: String -> String -> String -> String +replace what with = go + where + go [] = [] + go str@(s:ss) | what `isPrefixOf` str = with <> go (drop (length what) str) + | otherwise = s : go ss + defineAsmFun :: AsmCode tyAnn code => String -> tyAnn -> code -> Q [Dec] defineAsmFun name tyAnn asmCode = do addForeignSource LangAsm $ unlines [ ".global " <> asmName , asmName <> ":" - , codeToString tyAnn asmCode - , "jmp *(%rbp)" + , replace "RET_HASK" retToHask $ codeToString tyAnn asmCode + , retToHask ] funTy <- toTypeQ tyAnn let importedName = mkName asmName @@ -60,11 +73,17 @@ where name' = mkName name asmName = name <> "_unlifted" + retToHask = "jmp *(%rbp)" +getBSAddr :: BS.ByteString -> Ptr Word8 +getBSAddr bs = unsafeForeignPtrToPtr ptr `plusPtr` offset + where + (ptr, offset, _) = BS.toForeignPtr bs + mkFunD :: String -> Name -> Type -> Q Dec mkFunD funName importedName funTy = do argNames <- replicateM (countArgs funTy) $ newName "arg" - funAppE <- foldM f (VarE importedName) argNames + funAppE <- foldM f (VarE importedName) $ zip (VarE <$> argNames) (getArgs funTy) body <- case detectRetTuple funTy of Nothing -> [e| rebox $(pure funAppE) |] Just n -> do @@ -75,19 +94,30 @@ |] pure $ FunD (mkName funName) [Clause (VarP <$> argNames) (NormalB body) []] where - f acc argName = [e| $(pure acc) (unbox $(pure $ VarE argName)) |] + f acc (argName, argType) | argType == ConT ''BS.ByteString = [e| $(pure acc) + (unbox $ getBSAddr $(pure argName)) + (unbox $ BS.length $(pure argName)) + |] + | otherwise = [e| $(pure acc) (unbox $(pure argName)) |] unliftType :: Type -> Type -unliftType = transformBi unliftTuple . transformBi unliftBaseTy . transformBi unliftPtrs +unliftType = transformBi unliftTuple + . transformBi unliftBaseTy + . transformBi unliftPtrs + . transformBi unliftBS where unliftBaseTy x | x == ''Word = ''Word# | x == ''Int = ''Int# | x == ''Double = ''Double# | x == ''Float = ''Float# | otherwise = x + unliftPtrs (AppT (ConT name) _) | name == ''Ptr = ConT ''Addr# unliftPtrs x = x + unliftBS (AppT (AppT ArrowT (ConT bs)) rhs) | bs == ''BS.ByteString = unsafePerformIO $ runQ [t| Addr# -> Int# -> $(pure rhs) |] + unliftBS x = x + unliftTuple (TupleT n) = UnboxedTupleT n unliftTuple x = x @@ -96,3 +126,6 @@ detectRetTuple (AppT lhs _) = detectRetTuple lhs detectRetTuple (TupleT n) = Just n detectRetTuple _ = Nothing + +getArgs :: Type -> [Type] +getArgs ty = [ argTy | AppT ArrowT argTy <- universeBi ty ] diff --git a/src/Language/Asm/Inline/QQ.hs b/src/Language/Asm/Inline/QQ.hs --- a/src/Language/Asm/Inline/QQ.hs +++ b/src/Language/Asm/Inline/QQ.hs @@ -60,23 +60,28 @@ newtype RegName = RegName { regName :: String } deriving (Show, IsString) computeRegisters :: [(AsmVarName, AsmVarType)] -> Either String [(AsmVarName, RegName)] -computeRegisters vars = fst <$> foldM f ([], mempty) vars +computeRegisters vars = fst <$> foldM handleType ([], mempty) vars where - f (regNames, regCounts) (name, ty) = do - cat <- categorize ty - let idx = M.findWithDefault 0 cat regCounts + handleType (regNames, regCounts) (name, ty) = do + cats <- categorize name ty + foldM handleCats (regNames, regCounts) cats + + handleCats (regNames, regCounts) (name, cat) = do reg <- argIdxToReg cat idx pure ((name, reg) : regNames, M.insert cat (idx + 1) regCounts) + where + idx = M.findWithDefault 0 cat regCounts data VarTyCat = Integer | Other deriving (Eq, Ord, Show, Enum, Bounded) -categorize :: AsmVarType -> Either String VarTyCat -categorize (AsmVarType "Int") = pure Integer -categorize (AsmVarType "Word") = pure Integer -categorize (AsmVarType "Ptr") = pure Integer -categorize (AsmVarType "Float") = pure Other -categorize (AsmVarType "Double") = pure Other -categorize (AsmVarType s) = throwError $ "Unknown register type: " <> s +categorize :: AsmVarName -> AsmVarType -> Either String [(AsmVarName, VarTyCat)] +categorize name (AsmVarType "Int") = pure [(name, Integer)] +categorize name (AsmVarType "Word") = pure [(name, Integer)] +categorize name (AsmVarType "Ptr") = pure [(name, Integer)] +categorize name (AsmVarType "Float") = pure [(name, Other)] +categorize name (AsmVarType "Double") = pure [(name, Other)] +categorize name (AsmVarType "ByteString") = pure [(name <> ":ptr", Integer), (name <> ":len", Integer)] +categorize _ (AsmVarType s) = throwError $ "Unknown register type: " <> s argIdxToReg :: VarTyCat -> Int -> Either String RegName argIdxToReg Integer 0 = pure "rbx" @@ -111,7 +116,7 @@ Left err -> error err Right parsed -> [e| parsed |] -newtype AsmVarName = AsmVarName { varName :: String } deriving (Show, Eq, Lift) +newtype AsmVarName = AsmVarName { varName :: String } deriving (Show, Eq, Lift, Semigroup, IsString) newtype AsmVarType = AsmVarType { varType :: String } deriving (Show, Eq, Lift) data AsmQQType = AsmQQType diff --git a/test/Spec.hs b/test/Spec.hs --- a/test/Spec.hs +++ b/test/Spec.hs @@ -1,6 +1,8 @@ {-# LANGUAGE TemplateHaskell, QuasiQuotes #-} {-# LANGUAGE GHCForeignImportPrim, UnliftedFFITypes, UnboxedTuples #-} +import qualified Data.ByteString.Char8 as BS +import Data.ByteString(ByteString) import Foreign.Ptr import Test.Hspec import Test.QuickCheck @@ -69,6 +71,19 @@ xchg ${a}, ${b} |] + +defineAsmFun "lastChar" + [asmTy| (bs : ByteString) (def : Word) | (w : Word) |] + [asm| + test ${bs:len}, ${bs:len} + jz is_zero + movzbq -1(${bs:ptr},${bs:len}), ${w} + RET_HASK +is_zero: + mov ${def}, ${w} + |] + + main :: IO () main = hspec $ do describe "Works with Ints (the non-QQ version)" $ do @@ -93,6 +108,11 @@ in swapPtrs p1 p2 `shouldBe` (p2, p1) describe "Works on mixed types" $ it "timesTwoEverything" $ property $ \d n f w -> timesTwoEverything d n f w `shouldBe` (d + d, n * 2, f + f, w * 2) + describe "Works on ByteString" $ + it "lastChar" $ property $ \(ASCIIString str) def -> let bs = BS.pack str + in lastChar bs def `shouldBe` if BS.null bs + then def + else fromIntegral $ fromEnum $ BS.last bs intToPtr :: Int -> Ptr a intToPtr = intPtrToPtr . IntPtr