[commit: ghc] master: Fix memory leak from #12664 (e41b9c6)

[git at git.haskell.org]

Repository : ssh://git@git.haskell.org/ghc

On branch  : master
Link       : http://ghc.haskell.org/trac/ghc/changeset/e41b9c614984b63c4660018cecde682453e083e5/ghc

>---------------------------------------------------------------

commit e41b9c614984b63c4660018cecde682453e083e5
Author: Bartosz Nitka <niteria at gmail.com>
Date:   Thu Oct 6 05:40:24 2016 -0700

    Fix memory leak from #12664
    
    This fixes the leak with `setProgArgv`. The problem was
    that `setProgArgv` would not free the objects pointed
    to by `prog_argc`, `prog_argv` when the globals were
    changed resulting in a leak.
    
    The only strictly necessary change is in `rts/RtsFlags.c`, but
    the code in `System.Environment` was a bit confusing and not
    exception safe, so I refactored it.
    
    Test Plan: ./validate
    
    Reviewers: simonmar, ezyang, austin, hvr, bgamari, erikd
    
    Subscribers: thomie
    
    Differential Revision: https://phabricator.haskell.org/D2576
    
    GHC Trac Issues: #12664


>---------------------------------------------------------------

e41b9c614984b63c4660018cecde682453e083e5
 libraries/base/GHC/Foreign.hs                 | 18 ++++++++++++++++
 libraries/base/System/Environment.hs          | 30 +++++++++++----------------
 libraries/base/tests/IO/environment001.hs     |  4 ++++
 libraries/base/tests/IO/environment001.stdout |  2 ++
 rts/RtsFlags.c                                |  1 +
 5 files changed, 37 insertions(+), 18 deletions(-)

diff --git a/libraries/base/GHC/Foreign.hs b/libraries/base/GHC/Foreign.hs
index e8553d8..7d2f915 100644
--- a/libraries/base/GHC/Foreign.hs
+++ b/libraries/base/GHC/Foreign.hs
@@ -32,6 +32,7 @@ module GHC.Foreign (
     --
     withCString,
     withCStringLen,
+    withCStringsLen,
 
     charIsRepresentable,
   ) where
@@ -134,6 +135,23 @@ withCString enc s act = withEncodedCString enc True s $ \(cp, _sz) -> act cp
 withCStringLen         :: TextEncoding -> String -> (CStringLen -> IO a) -> IO a
 withCStringLen enc = withEncodedCString enc False
 
+-- | Marshal a list of Haskell strings into an array of NUL terminated C strings
+-- using temporary storage.
+--
+-- * the Haskell strings may /not/ contain any NUL characters
+--
+-- * the memory is freed when the subcomputation terminates (either
+--   normally or via an exception), so the pointer to the temporary
+--   storage must /not/ be used after this.
+--
+withCStringsLen :: TextEncoding
+                -> [String]
+                -> (Int -> Ptr CString -> IO a)
+                -> IO a
+withCStringsLen enc strs f = go [] strs
+  where
+  go cs (s:ss) = withCString enc s $ \c -> go (c:cs) ss
+  go cs [] = withArrayLen (reverse cs) f
 
 -- | Determines whether a character can be accurately encoded in a 'CString'.
 --
diff --git a/libraries/base/System/Environment.hs b/libraries/base/System/Environment.hs
index 242845a..d8b3e03 100644
--- a/libraries/base/System/Environment.hs
+++ b/libraries/base/System/Environment.hs
@@ -32,12 +32,14 @@ module System.Environment
 import Foreign
 import Foreign.C
 import System.IO.Error (mkIOError)
-import Control.Exception.Base (bracket, throwIO)
+import Control.Exception.Base (bracket_, throwIO)
+#ifdef mingw32_HOST_OS
+import Control.Exception.Base (bracket)
+#endif
 -- import GHC.IO
 import GHC.IO.Exception
 import GHC.IO.Encoding (getFileSystemEncoding)
 import qualified GHC.Foreign as GHC
-import Data.List
 import Control.Monad
 #ifdef mingw32_HOST_OS
 import GHC.Environment
@@ -369,25 +371,17 @@ withProgArgv :: [String] -> IO a -> IO a
 withProgArgv new_args act = do
   pName <- System.Environment.getProgName
   existing_args <- System.Environment.getArgs
-  bracket (setProgArgv new_args)
-          (\argv -> do _ <- setProgArgv (pName:existing_args)
-                       freeProgArgv argv)
-          (const act)
-
-freeProgArgv :: Ptr CString -> IO ()
-freeProgArgv argv = do
-  size <- lengthArray0 nullPtr argv
-  sequence_ [ peek (argv `advancePtr` i) >>= free
-            | i <- [size - 1, size - 2 .. 0]]
-  free argv
-
-setProgArgv :: [String] -> IO (Ptr CString)
+  bracket_ (setProgArgv new_args)
+           (setProgArgv (pName:existing_args))
+           act
+
+setProgArgv :: [String] -> IO ()
 setProgArgv argv = do
   enc <- getFileSystemEncoding
-  vs <- mapM (GHC.newCString enc) argv >>= newArray0 nullPtr
-  c_setProgArgv (genericLength argv) vs
-  return vs
+  GHC.withCStringsLen enc argv $ \len css ->
+    c_setProgArgv (fromIntegral len) css
 
+-- setProgArgv copies the arguments
 foreign import ccall unsafe "setProgArgv"
   c_setProgArgv  :: CInt -> Ptr CString -> IO ()
 
diff --git a/libraries/base/tests/IO/environment001.hs b/libraries/base/tests/IO/environment001.hs
index 11d7912..1d7a5c1 100644
--- a/libraries/base/tests/IO/environment001.hs
+++ b/libraries/base/tests/IO/environment001.hs
@@ -14,3 +14,7 @@ main = do
     [arg1] <- withArgs ["你好!"] getArgs
     putStrLn arg1
     putStrLn ("Test 3: " ++ show (length arg1))
+
+    args2 <- withArgs ["a", "b"] getArgs
+    print args2
+    putStrLn ("Test 4: " ++ show (length args2))
diff --git a/libraries/base/tests/IO/environment001.stdout b/libraries/base/tests/IO/environment001.stdout
index 2434d0c..2d32a83 100644
--- a/libraries/base/tests/IO/environment001.stdout
+++ b/libraries/base/tests/IO/environment001.stdout
@@ -4,3 +4,5 @@ Test 1: 3
 Test 2: 1
 你好!
 Test 3: 3
+["a","b"]
+Test 4: 2
diff --git a/rts/RtsFlags.c b/rts/RtsFlags.c
index c994a0c..4bd544e 100644
--- a/rts/RtsFlags.c
+++ b/rts/RtsFlags.c
@@ -1943,6 +1943,7 @@ getProgArgv(int *argc, char **argv[])
 void
 setProgArgv(int argc, char *argv[])
 {
+    freeArgv(prog_argc,prog_argv);
     prog_argc = argc;
     prog_argv = copyArgv(argc,argv);
     setProgName(prog_argv);