diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,5 +1,10 @@ # Changelog for inline-asm +## v0.3.0.0 + +* Support named arguments for doubles and floats as well. +* Support passing Ptr values. + ## v0.2.0.0 * Support named arguments. diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -17,11 +17,12 @@ main = print $ timesTwo 21 ``` -There is also an alternative notation allowing named arguments: +There is also an alternative notation allowing named arguments to avoid remembering +which arguments are passed in which registers: ```haskell defineAsmFun "swap2p1" - [t| Int -> Int -> (Int, Int) |] - [asm| a b | + [asmTy| (a : Int) (b : Int) | (_ : Int) (_ : Int)] + [asm| xchg ${a}, ${b} add $1, ${b} |] diff --git a/app/Main.hs b/app/Main.hs --- a/app/Main.hs +++ b/app/Main.hs @@ -9,6 +9,14 @@ defineAsmFun "swap" [t| Int -> Int -> (Int, Int) |] "xchg %rbx, %r14" defineAsmFun "swap2p1" + [asmTy| (a : Int) (b : Int) | (a : Int) (b : Int) |] + [asm| + xchg ${a}, ${b} + add $1, ${b} + |] + +{- +defineAsmFun "swap2p1" [t| Int -> Int -> (Int, Int) |] [asm| a b | xchg ${a}, ${b} @@ -17,8 +25,10 @@ defineAsmFun "testInt" [t| Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int |] "int $3" defineAsmFun "testDouble" [t| Double -> Double -> Double -> Double -> Double -> Double -> Double -> Double |] "int $3" +-} main :: IO () main = do - print $ testInt 0 1 2 3 4 5 6 - print $ testDouble 1 1 0 0 0 0 1 + print $ swap2p1 2 4 + --print $ testInt 0 1 2 3 4 5 6 + --print $ testDouble 1 1 0 0 0 0 1 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: b58246a639e949912f040d75472b9a114b43d99ceb097b6ff9df53c0c370b8d4 +-- hash: 07f84f8e6bb0bdb11a97883ebf8eb73d17aeafe6f29f05996646e7487b23583f name: inline-asm -version: 0.2.0.0 +version: 0.2.1.0 synopsis: Inline some Assembly in ur Haskell! description: Please see the README on GitHub at category: FFI @@ -40,8 +40,10 @@ ghc-options: -Wall build-depends: base >=4.7 && <5 + , containers , either , ghc-prim + , megaparsec , mtl , template-haskell >=2.15.0.0 , uniplate @@ -56,9 +58,11 @@ ghc-options: -Wall -threaded -rtsopts -with-rtsopts=-N build-depends: base >=4.7 && <5 + , containers , either , ghc-prim , inline-asm + , megaparsec , mtl , template-haskell >=2.15.0.0 , uniplate @@ -75,10 +79,12 @@ build-depends: QuickCheck , base >=4.7 && <5 + , containers , either , ghc-prim , hspec , inline-asm + , megaparsec , mtl , template-haskell >=2.15.0.0 , uniplate 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 @@ -7,7 +7,9 @@ import Control.Monad import Data.Generics.Uniplate.Data +import Foreign.Ptr import GHC.Prim +import GHC.Ptr import GHC.Types hiding (Type) import Language.Haskell.TH import Language.Haskell.TH.Syntax @@ -35,17 +37,18 @@ unbox (F# f) = f rebox = F# -defineAsmFun :: AsmCode c => String -> Q Type -> c -> Q [Dec] -defineAsmFun name funTyQ asmCode = do +instance AsmArg (Ptr a) 'AddrRep Addr# where + unbox (Ptr p) = p + rebox = Ptr + +defineAsmFun :: AsmCode tyAnn code => String -> tyAnn -> code -> Q [Dec] +defineAsmFun name tyAnn asmCode = do addForeignSource LangAsm $ unlines [ ".global " <> asmName , asmName <> ":" - , codeToString asmCode + , codeToString tyAnn asmCode , "jmp *(%rbp)" ] - funTy <- funTyQ - case validateCode funTy asmCode of - Right () -> pure () - Left err -> error err + funTy <- toTypeQ tyAnn let importedName = mkName asmName wrapperFunD <- mkFunD name importedName funTy pure @@ -75,21 +78,21 @@ f acc argName = [e| $(pure acc) (unbox $(pure $ VarE argName)) |] unliftType :: Type -> Type -unliftType = transformBi unliftTuple . transformBi unliftBaseTy +unliftType = transformBi unliftTuple . transformBi unliftBaseTy . transformBi unliftPtrs 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 + unliftTuple (TupleT n) = UnboxedTupleT n unliftTuple x = x --- This doesn't check if this is indeed a return type, --- but since we are not going to support argument tuples (and we'll add a check about that later), --- it should be fine. detectRetTuple :: Type -> Maybe Int -detectRetTuple ty | [TupleT n] <- tuples = Just n - | otherwise = Nothing - where - tuples = [ t | t@(TupleT _) <- universeBi ty] +detectRetTuple (AppT (AppT ArrowT _) rhs) = detectRetTuple rhs +detectRetTuple (AppT lhs _) = detectRetTuple lhs +detectRetTuple (TupleT n) = Just n +detectRetTuple _ = Nothing diff --git a/src/Language/Asm/Inline/AsmCode.hs b/src/Language/Asm/Inline/AsmCode.hs --- a/src/Language/Asm/Inline/AsmCode.hs +++ b/src/Language/Asm/Inline/AsmCode.hs @@ -1,13 +1,13 @@ -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FlexibleInstances, FunctionalDependencies #-} module Language.Asm.Inline.AsmCode where import Language.Haskell.TH.Syntax -class AsmCode c where - codeToString :: c -> String - validateCode :: Type -> c -> Either String () +class AsmCode tyAnn code | code -> tyAnn, tyAnn -> code where + codeToString :: tyAnn -> code -> String + toTypeQ :: tyAnn -> Q Type -instance AsmCode String where - codeToString = id - validateCode _ _ = pure () +instance AsmCode (Q Type) String where + codeToString _ = id + toTypeQ = id 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 @@ -1,78 +1,178 @@ {-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE DeriveLift #-} -{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE FlexibleContexts, MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings, RecordWildCards, MultiWayIf #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE ScopedTypeVariables #-} -module Language.Asm.Inline.QQ(asm) where +module Language.Asm.Inline.QQ(asm, asmTy) where -import Control.Arrow +import qualified Data.Map as M import Control.Monad.Except +import Data.Bifunctor +import Data.Char import Data.Either.Combinators -import Data.List +import Data.Foldable +import Data.String +import Data.Void +import Foreign.Ptr import Language.Haskell.TH import Language.Haskell.TH.Quote import Language.Haskell.TH.Syntax +import Text.Megaparsec +import Text.Megaparsec.Char +import qualified Text.Megaparsec.Char.Lexer as ML import Language.Asm.Inline.AsmCode -import Language.Asm.Inline.Util -instance AsmCode AsmQQParsed where - codeToString = asmBody - validateCode ty code = - check "arguments count mismatch" $ argsCount code == countArgs ty - where - check _ True = pure () - check str False = throwError $ "Type error: " <> str +instance AsmCode AsmQQType AsmQQCode where + codeToString ty code = case substituteArgs ty code of + Left e -> error e + Right s -> s + toTypeQ = unreflectTy asm :: QuasiQuoter -asm = QuasiQuoter { quoteExp = asmQE, quotePat = unsupported, quoteType = unsupported, quoteDec = unsupported } - where - unsupported = const $ error "Unsupported quasiquotation type" +asm = expQQ asmQE asmQE :: String -> Q Exp -asmQE p = case parseAsmQQ p of - Left err -> error err - Right parsed -> [e| parsed |] +asmQE p = [e| AsmQQCode p |] -data AsmQQParsed = AsmQQParsed - { argsCount :: Int - , asmBody :: String - } deriving (Show, Lift) +newtype AsmQQCode = AsmQQCode { asmCode :: String } -parseAsmQQ :: String -> Either String AsmQQParsed -parseAsmQQ = findSplitter - >=> (pure . first words) - >=> substituteArgs +substituteArgs :: AsmQQType -> AsmQQCode -> Either String String +substituteArgs AsmQQType { .. } AsmQQCode { .. } = do + argRegs <- computeRegisters args + retRegs <- computeRegisters rets + go' argRegs retRegs asmCode where - findSplitter p = case break (== '|') p of - (vars, '|' : body) -> pure (vars, body) - _ -> throwError "Unable to find variable section separator" + go' argRegs retRegs = go + where + go ('$' : '{' : rest) + | (argStr, '}' : rest') <- break (== '}') rest + , not $ null argStr = do + let arg = AsmVarName $ trim argStr + RegName reg <- maybeToRight ("Unknown argument: `" <> show arg <> "`") $ msum [lookup arg argRegs, lookup arg retRegs] + (('%' : reg) <>) <$> go rest' + | otherwise = throwError $ "Unable to parse argument: " <> take 20 rest <> "..." + go (x : xs) = (x :) <$> go xs + go [] = pure [] -substituteArgs :: ([String], String) -> Either String AsmQQParsed -substituteArgs (args, contents) = AsmQQParsed (length args) <$> go contents +newtype RegName = RegName { regName :: String } deriving (Show, IsString) + +computeRegisters :: [(AsmVarName, AsmVarType)] -> Either String [(AsmVarName, RegName)] +computeRegisters vars = fst <$> foldM f ([], mempty) vars where - go ('$' : '{' : rest) - | (arg, '}' : rest') <- break (== '}') rest = do - idx <- if retPrefix `isPrefixOf` arg - then pure $ read (drop (length retPrefix) arg) - else maybeToRight ("Unknown argument: `" <> trim arg <> "`") $ elemIndex (trim arg) args - reg <- argIdxToReg idx - (('%' : reg) <>) <$> go rest' - | otherwise = throwError $ "Unable to parse argument: " <> take 20 rest <> "..." - go (x : xs) = (x :) <$> go xs - go [] = pure [] + f (regNames, regCounts) (name, ty) = do + cat <- categorize ty + let idx = M.findWithDefault 0 cat regCounts + reg <- argIdxToReg cat idx + pure ((name, reg) : regNames, M.insert cat (idx + 1) regCounts) - retPrefix = "ret" +data VarTyCat = Integer | Other deriving (Eq, Ord, Show, Enum, Bounded) -argIdxToReg :: Int -> Either String String -argIdxToReg 0 = pure "rbx" -argIdxToReg 1 = pure "r14" -argIdxToReg 2 = pure "rsi" -argIdxToReg 3 = pure "rdi" -argIdxToReg 4 = pure "r8" -argIdxToReg 5 = pure "r9" -argIdxToReg n = throwError $ "Unsupported register index: " <> show n +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 +argIdxToReg :: VarTyCat -> Int -> Either String RegName +argIdxToReg Integer 0 = pure "rbx" +argIdxToReg Integer 1 = pure "r14" +argIdxToReg Integer 2 = pure "rsi" +argIdxToReg Integer 3 = pure "rdi" +argIdxToReg Integer 4 = pure "r8" +argIdxToReg Integer 5 = pure "r9" +argIdxToReg Other n | n >= 0 && n <= 6 = pure $ RegName $ "xmm" <> show (n + 1) +argIdxToReg _ n = throwError $ "Unsupported register index: " <> show n + trim :: String -> String trim = pass . pass where pass = reverse . dropWhile (== ' ') + +findSplitter :: String -> Either String (String, String) +findSplitter p = case break (== '|') p of + (vars, '|' : body) -> pure (vars, body) + _ -> throwError "Unable to find variable section separator" + +expQQ :: (String -> Q Exp) -> QuasiQuoter +expQQ qq = QuasiQuoter { quoteExp = qq, quotePat = unsupported, quoteType = unsupported, quoteDec = unsupported } + where + unsupported = const $ error "Unsupported quasiquotation type" + +asmTy :: QuasiQuoter +asmTy = expQQ asmTyQE + +asmTyQE :: String -> Q Exp +asmTyQE str = case parseAsmTyQQ str of + Left err -> error err + Right parsed -> [e| parsed |] + +newtype AsmVarName = AsmVarName { varName :: String } deriving (Show, Eq, Lift) +newtype AsmVarType = AsmVarType { varType :: String } deriving (Show, Eq, Lift) + +data AsmQQType = AsmQQType + { args :: [(AsmVarName, AsmVarType)] + , rets :: [(AsmVarName, AsmVarType)] + } deriving (Show, Lift) + +parseAsmTyQQ :: String -> Either String AsmQQType +parseAsmTyQQ str = do + (inputStr, outputStr) <- findSplitter str + args <- first showParseError $ runParser (parseInTypes <* eof) "" inputStr + rets <- first showParseError $ runParser (parseInTypes <* eof) "" outputStr + pure AsmQQType { .. } + where + showParseError = errorBundlePretty :: ParseErrorBundle String Void -> String + +parseInTypes :: forall m e. MonadParsec e String m => m [(AsmVarName, AsmVarType)] +parseInTypes = space *> many parseType + where + parseType = do + void $ lexeme $ string "(" + name <- lexeme $ parseWFirst letterChar <|> string "_" + void $ lexeme $ string ":" + ty <- lexeme $ parseWFirst upperChar + void $ takeWhileP Nothing (/= ')') + void $ lexeme $ string ")" + pure (AsmVarName name, AsmVarType ty) + + parseWFirst :: m Char -> m String + parseWFirst p = do + firstLetter <- p + rest <- takeWhileP (Just "variable") isAlphaNum + pure $ firstLetter : rest + + lexeme = ML.lexeme $ ML.space space1 empty empty + +unreflectTy :: AsmQQType -> Q Type +unreflectTy AsmQQType { .. } = do + retTy <- unreflectRetTy rets + maybeArgTyNames <- lookupTyNames args + case maybeArgTyNames of + Left err -> error err + Right argTyNames -> foldrM argFolder retTy argTyNames + where + argFolder argName funAcc | argName == ''Ptr = [t| Ptr () -> $(pure funAcc) |] + | otherwise = [t| $(pure $ ConT argName) -> $(pure funAcc) |] + +unreflectRetTy :: [(AsmVarName, AsmVarType)] -> Q Type +unreflectRetTy [] = [t| () |] +unreflectRetTy rets = do + maybeRetTyNames <- lookupTyNames rets + case maybeRetTyNames of + Left err -> error err + Right [tyName] -> if | tyName == ''Ptr -> [t| Ptr () |] + | otherwise -> pure $ ConT tyName + Right retNames -> pure $ foldl retFolder (TupleT $ length retNames) retNames + where + retFolder tupAcc ret | ret == ''Ptr = tupAcc `AppT` (ConT ret `AppT` TupleT 0) + | otherwise = tupAcc `AppT` ConT ret + +lookupTyNames :: [(AsmVarName, AsmVarType)] -> Q (Either String [Name]) +lookupTyNames = fmap sequence . mapM f + where + f (name, ty) = maybeToRight ("Unable to lookup type " <> show ty <> " for var " <> show name) <$> lookupTypeName (varType ty) diff --git a/test/Spec.hs b/test/Spec.hs --- a/test/Spec.hs +++ b/test/Spec.hs @@ -1,19 +1,98 @@ -{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TemplateHaskell, QuasiQuotes #-} {-# LANGUAGE GHCForeignImportPrim, UnliftedFFITypes, UnboxedTuples #-} +import Foreign.Ptr import Test.Hspec import Test.QuickCheck import Language.Asm.Inline +import Language.Asm.Inline.QQ -defineAsmFun "timesTwo" [t| Int -> Int |] "add %rbx, %rbx" -defineAsmFun "plusWord" [t| Int -> Int -> Int |] "add %r14, %rbx" -defineAsmFun "swap" [t| Int -> Int -> (Int, Int) |] "xchg %rbx, %r14" +defineAsmFun "timesTwoInt" [t| Int -> Int |] "add %rbx, %rbx" +defineAsmFun "plusInt" [t| Int -> Int -> Int |] "add %r14, %rbx" +defineAsmFun "swapInts" [t| Int -> Int -> (Int, Int) |] "xchg %rbx, %r14" + +defineAsmFun "timesTwoIntQQ" + [asmTy| (a : Int) | (_ : Int) |] + [asm| add ${a}, ${a} |] + +defineAsmFun "plusIntQQ" + [asmTy| (a : Int) (b : Int) | (_ : Int) |] + [asm| add ${b}, ${a} |] + +defineAsmFun "swap2p1QQ" + [asmTy| (a : Int) (b : Int) | (a : Int) (b : Int) |] + [asm| + xchg ${a}, ${b} + add $1, ${b} + |] + + +defineAsmFun "timesTwoFloatQQ" + [asmTy| (a : Float) | (_ : Float) |] + [asm| addss ${a}, ${a} |] + +defineAsmFun "plusFloatQQ" + [asmTy| (a : Float) (b : Float) | (_ : Float) |] + [asm| addss ${b}, ${a} |] + + +defineAsmFun "timesTwoDoubleQQ" + [asmTy| (a : Double) | (_ : Double) |] + [asm| addsd ${a}, ${a} |] + +defineAsmFun "plusDoubleQQ" + [asmTy| (a : Double) (b : Double) | (_ : Double) |] + [asm| addsd ${b}, ${a} |] + + +defineAsmFun "timesTwoEverything" + [asmTy| (d : Double) (n : Int) (f : Float) (w : Word) | (_ : Double) (_ : Int) (_ : Float) (_ : Word) |] + [asm| + addsd ${d}, ${d} + addss ${f}, ${f} + add ${n}, ${n} + add ${w}, ${w} + |] + + +defineAsmFun "addPtr" + [asmTy| (ptr : Ptr Int) (shift : Int) | (_ : Ptr Int) |] + [asm| + add ${shift}, ${ptr} + |] + +defineAsmFun "swapPtrs" + [asmTy| (a : Ptr Int) (b : Ptr Int) | (_ : Ptr Int) (_ : Ptr Int) |] + [asm| + xchg ${a}, ${b} + |] + main :: IO () main = hspec $ do - describe "defineAsmFun" $ do - it "timesTwo" $ property $ \num -> timesTwo num `shouldBe` num * 2 - it "plusWord" $ property $ \n1 n2 -> plusWord n1 n2 `shouldBe` n1 + n2 - describe "Returning tuples" $ - it "swap" $ property $ \n1 n2 -> swap n1 n2 `shouldBe` (n2, n1) + describe "Works with Ints (the non-QQ version)" $ do + it "timesTwo" $ property $ \num -> timesTwoInt num `shouldBe` num * 2 + it "plusWord" $ property $ \n1 n2 -> plusInt n1 n2 `shouldBe` n1 + n2 + it "swap returns a tuple properly" $ property $ \n1 n2 -> swapInts n1 n2 `shouldBe` (n2, n1) + describe "Works on Ints" $ do + it "timesTwo" $ property $ \num -> timesTwoIntQQ num `shouldBe` num * 2 + it "plusWord" $ property $ \n1 n2 -> plusIntQQ n1 n2 `shouldBe` n1 + n2 + it "swap returns a tuple properly" $ property $ \n1 n2 -> swap2p1QQ n1 n2 `shouldBe` (n2, n1 + 1) + describe "Works on Floats" $ do + it "timesTwo" $ property $ \num -> timesTwoFloatQQ num `shouldBe` num * 2 + it "plusWord" $ property $ \n1 n2 -> plusFloatQQ n1 n2 `shouldBe` n1 + n2 + describe "Works on Doubles" $ do + it "timesTwo" $ property $ \num -> timesTwoDoubleQQ num `shouldBe` num * 2 + it "plusWord" $ property $ \n1 n2 -> plusDoubleQQ n1 n2 `shouldBe` n1 + n2 + describe "Works on Ptrs" $ do + it "addPtr" $ property $ \int num -> let ptr = intToPtr int + in addPtr ptr num `shouldBe` (ptr `plusPtr` num) + it "swap returns a tuple properly" $ property $ \n1 n2 -> let p1 = intToPtr n1 + p2 = intToPtr n2 + 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) + +intToPtr :: Int -> Ptr a +intToPtr = intPtrToPtr . IntPtr