Initial

6 files changed, 130 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..a3ac1fc
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+dist-newstyle/
+cabal.project.local
diff --git a/init-once.cabal b/init-once.cabal
new file mode 100644
index 0000000..a02dab5
--- /dev/null
+++ b/init-once.cabal
@@ -0,0 +1,29 @@
+cabal-version:   3.0
+name:            init-once
+version:         0.1.0.0
+license:         BSD-3-Clause
+author:          Tom Smeding
+maintainer:      Tom Smeding <tom@tomsmeding.com>
+build-type:      Simple
+
+library
+  exposed-modules:
+    InitOnce
+  other-modules:
+    InitOnce.TheMap
+    InitOnce.Types
+  build-depends:
+    base >= 4.17,
+    containers,
+    template-haskell
+  hs-source-dirs: src
+  default-language: Haskell2010
+  ghc-options: -Wall
+
+test-suite hs-init-once-test
+  type: exitcode-stdio-1.0
+  main-is: Main.hs
+  build-depends: base, init-once
+  hs-source-dirs: test
+  default-language: Haskell2010
+  ghc-options: -Wall
diff --git a/src/InitOnce.hs b/src/InitOnce.hs
new file mode 100644
index 0000000..60ea435
--- /dev/null
+++ b/src/InitOnce.hs
@@ -0,0 +1,18 @@
+{-# LANGUAGE TemplateHaskellQuotes #-}
+module InitOnce (once) where
+
+import Language.Haskell.TH (bindCode, CodeQ)
+import Type.Reflection (Typeable, typeOf)
+
+import InitOnce.TheMap (onceInsert)
+import InitOnce.Types (Location, location', TypedAny(..), getTypedAny)
+
+
+-- | The IO action is executed with asynchronous exceptions masked (using
+-- 'Control.Concurrent.MVar.modifyMVarMasked') to ensure that initialisation
+-- really happens once.
+once :: Typeable a => CodeQ (IO a) -> CodeQ a
+once code = bindCode location' $ \loc -> [|| typedOnceInsert loc $$(code) ||]
+
+typedOnceInsert :: Typeable a => Location -> IO a -> a
+typedOnceInsert loc act = getTypedAny (onceInsert loc ((\x -> TypedAny (typeOf x) x) <$> act))
diff --git a/src/InitOnce/TheMap.hs b/src/InitOnce/TheMap.hs
new file mode 100644
index 0000000..5c30794
--- /dev/null
+++ b/src/InitOnce/TheMap.hs
@@ -0,0 +1,32 @@
+-- | This is the only module that uses unsafePerformIO; to make sure it's safe,
+-- we disable all the known optimisations in GHC that can potentially mess
+-- things up: <https://discourse.haskell.org/t/can-noinline-fail-to-prevent-inlining/14324>
+-- Also, since there is nothing else in this module, there isn't really very
+-- much for GHC /to/ mess up, which is the point.
+{-# OPTIONS -fno-cse -fno-full-laziness -fno-state-hack #-}
+module InitOnce.TheMap (onceInsert) where
+
+import Control.Concurrent.MVar (modifyMVarMasked, newMVar, MVar)
+import Data.Map.Strict (Map)
+import qualified Data.Map.Strict as Map (empty, lookup, insert)
+import System.IO.Unsafe (unsafePerformIO)
+
+import InitOnce.Types (Location, TypedAny)
+
+
+{-# NOINLINE theMap #-}
+theMap :: MVar (Map Location TypedAny)
+theMap = unsafePerformIO (newMVar Map.empty)
+
+-- | The IO action is executed with asynchronous exceptions masked (using
+-- 'modifyMVarMasked').
+{-# NOINLINE onceInsert #-}
+onceInsert :: Location -> IO TypedAny -> TypedAny
+onceInsert key mkVal =
+  unsafePerformIO $
+    modifyMVarMasked theMap $ \mp ->
+      case Map.lookup key mp of
+        Nothing -> do
+          v <- mkVal
+          return (Map.insert key v mp, v)
+        Just v  -> return (mp, v)
diff --git a/src/InitOnce/Types.hs b/src/InitOnce/Types.hs
new file mode 100644
index 0000000..05d5c9c
--- /dev/null
+++ b/src/InitOnce/Types.hs
@@ -0,0 +1,30 @@
+{-# LANGUAGE DeriveLift #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+module InitOnce.Types (Location, location', TypedAny(TypedAny), getTypedAny) where
+
+import Language.Haskell.TH (Q, location, Loc(loc_start, loc_package, loc_module))
+import Language.Haskell.TH.Syntax (CharPos, Lift)
+import Type.Reflection (Typeable, TypeRep, (:~~:)(HRefl), pattern TypeRep, eqTypeRep)
+
+
+data Location = Location
+    String   -- package
+    String   -- module
+    CharPos  -- position in source file (line)
+  deriving (Eq, Ord, Lift)
+
+location' :: Q Location
+location' = locationFromLoc <$> location
+  where locationFromLoc loc = Location (loc_package loc) (loc_module loc) (loc_start loc)
+
+data TypedAny = forall a. TypedAny (TypeRep a) a
+
+getTypedAny :: forall a. Typeable a => TypedAny -> a
+getTypedAny (TypedAny rep x) =
+  case eqTypeRep (TypeRep @a) rep of
+    Just HRefl -> x
+    Nothing ->
+      error $ "InitOnce: Type mismatch: requested " ++ show (TypeRep @a) ++ ", got " ++ show rep
diff --git a/test/Main.hs b/test/Main.hs
new file mode 100644
index 0000000..0f21a29
--- /dev/null
+++ b/test/Main.hs
@@ -0,0 +1,19 @@
+{-# LANGUAGE TemplateHaskell #-}
+module Main where
+
+import Data.IORef (IORef, atomicModifyIORef', newIORef)
+
+import InitOnce
+
+
+globalRef :: IORef Int
+globalRef = $$(once [|| newIORef 0 ||])
+
+{-# NOINLINE foo #-}
+foo :: IO String
+foo = atomicModifyIORef' globalRef (\i -> (i + 1, show i))
+
+main :: IO ()
+main = do
+  putStrLn =<< foo
+  putStrLn =<< foo