{-# LANGUAGE NoMonomorphismRestriction, Rank2Types #-}

-- Code generation with code movements via reference cells

module Anaphora where

import TSCore

import Data.IORef
import Control.Applicative
import Control.Monad (liftM2)

-- We can use STRef, IORef or any other reference cells
-- We chose IORef, as most undisciplined and so posing
-- the most danger for generating bad code
-- We could have used STRef


-- For the sake of abstraction, we define the following three
-- operations for working with mutable variables

with_ref :: (IORef (Maybe b) -> J IO repr a) -> J IO repr a
with_ref f = J (newIORef Nothing >>= unJ . f)

get_ref :: IORef (Maybe (J IO repr b)) -> J IO repr b
get_ref r = J $ do
  Just v <- readIORef r
  unJ v

-- Store the code value, and return it too
note_ref :: IORef (Maybe (J IO repr a)) -> J IO repr a -> J IO repr a
note_ref r m = J (writeIORef r (Just m) >> unJ m)


an1 = with_ref (\r -> (note_ref r (int 1)) +: get_ref r)

-- Store open code
an2 = lam (\x -> with_ref (\r -> (note_ref r (var x)) +: get_ref r))

-- store open code and let the effect cross the binder
an3 = lam (\x -> with_ref (\r -> 
          lam (\y -> weaken (note_ref r (var x)) +: var y) $$ get_ref r))

-- Attempting to leak a variable
{-
an4 = lam (\x -> with_ref (\r -> 
          lam (\y -> (note_ref r (var y)) +: var y) $$ get_ref r))

    Inferred type is less polymorphic than expected
      Quantified type variable `s' is mentioned in the environment:
        r :: IORef
               (Maybe (J IO (HV (H repr s Int, (H repr s1 a, h)) repr) Int))
          (bound at /home/oleg/papers/MetaFX/hybrid/Anaphora.hs:45:28)
    In the first argument of `lam', namely
        `(\ y -> (note_ref r (var y)) +: var y)'
-}

an1r = fmap (2 == ) $ runR an1

an1c = fmap ("(GHC.Num.+) 1 1" ==) $ runC an1


an2r = fmap (42 == ) . fmap ($ 21) $ runR an2

an2c = fmap ("\\x_0 -> (GHC.Num.+) x_0 x_0" ==) $ runC an2

an3r = fmap (42 == ) . fmap ($ 21) $ runR an3

an3c = fmap ("\\x_0 -> (\\x_1 -> (GHC.Num.+) x_0 x_1) x_0" ==) $ runC an3


-- Assertion insertion
-- A simpler version of the example from our PEPM09 paper

guarded_div1 :: (Applicative m, SSym repr, SymDIV repr, AssertPos repr) =>
     J m (HV h repr) Int -> J m (HV h repr) Int -> J m (HV h repr) Int
guarded_div1 x y = assertPos y (x /: y)

-- A supposedly complex computation
complex_exp :: (Applicative m, SSym repr) => J m (HV h repr) Int
complex_exp = int 100 *: int 200


exdiv1 = lam (\y -> complex_exp +: guarded_div1 (int 10) (var y))

exdiv1c = 
 "\\x_0 -> (GHC.Num.+) ((GHC.Num.*) 100 200) "++
 "(GHC.Base.assert (x_0 GHC.Classes.> 0) (GHC.Real.div 10 x_0))"
 == runCI exdiv1


-- Assertion-insertion locus
assert_locus :: (IORef (J IO repr a -> J IO repr a) -> J IO repr a) 
	        -> J IO repr a
assert_locus m = J $ do
  assert_code_ref <- newIORef id
  mv <- unJ $ m assert_code_ref
  transformer <- readIORef assert_code_ref
  unJ $ transformer (J (return mv))

add_assert :: IORef (a -> a) -> (a->a)
     -> J IO repr b
     -> J IO repr b
add_assert locus transformer m = J $ do
  modifyIORef locus ( . transformer)
  unJ m

guarded_div2 :: (SSym repr, AssertPos repr, SymDIV repr) =>
     IORef (J IO (HV h repr) a -> J IO (HV h repr) a)
     -> J IO (HV h repr) Int
     -> J IO (HV h repr) Int
     -> J IO (HV h repr) Int
guarded_div2 locus x y = 
    add_assert locus (assertPos y) $ x /: y

exdiv2 = lam (\y -> assert_locus $ \locus ->
	     complex_exp +: guarded_div2 locus (int 10) (var y))

exdiv2c = 
    fmap ("\\x_0 -> GHC.Base.assert (x_0 GHC.Classes.> 0) "++
	  "((GHC.Num.+) ((GHC.Num.*) 100 200) (GHC.Real.div 10 x_0))" ==) $
    runC exdiv2

guarded_div3 :: (SSym repr, AssertPos repr, Extends h h1, SymDIV repr) =>
     IORef (J IO (HV h repr) a -> J IO (HV h repr) a)
     -> J IO (HV h1 repr) Int
     -> J IO (HV h repr) Int
     -> J IO (HV h1 repr) Int
guarded_div3 locus x y =
    add_assert locus (assertPos y) $
    x /: weakens y

exdiv2' = lam (\y -> assert_locus $ \locus ->
	     complex_exp +: 
	       guarded_div3 locus (int 10) (var y))

exdiv2c' = 
    fmap ("\\x_0 -> GHC.Base.assert (x_0 GHC.Classes.> 0) "++
	  "((GHC.Num.+) ((GHC.Num.*) 100 200) (GHC.Real.div 10 x_0))" ==) $
    runC exdiv2'

exdiv3 = lam (\y -> assert_locus $ \locus ->
	   lam (\x ->
	     complex_exp +: 
	       guarded_div3 locus (var x) (var y)))
	      
exdiv3c = 
    fmap ("\\x_0 -> GHC.Base.assert (x_0 GHC.Classes.> 0) "++
	 "(\\x_1 -> (GHC.Num.+) ((GHC.Num.*) 100 200) (GHC.Real.div x_1 x_0))"
	  ==) $
    runC exdiv3


-- If we make a mistake and attempt to leak a binding
{-
exdiv3f = lam (\y -> assert_locus $ \locus ->
	   lam (\x ->
	     complex_exp +: 
	       guarded_div3 locus (var y) (var x)))

    Inferred type is less polymorphic than expected
      Quantified type variable `s' is mentioned in the environment:
        locus :: IORef
                   (J IO (HV (H repr s Int, h) repr) a
                    -> J IO (HV (H repr s Int, h) repr) a)
          (bound at /home/oleg/papers/MetaFX/hybrid/Anaphora.hs:155:37)
    In the first argument of `lam', namely
        `(\ x -> complex_exp +: guarded_div3 locus (var y) (var x))'
-}


-- Moving from differently nested environments

exdiv4 = lam (\y -> assert_locus $ \locus ->
	   lam (\x ->
	     complex_exp +: 
	       guarded_div3 locus (var x) (var y)) $$
	   complex_exp +: guarded_div3 locus (int 5) ((var y) +: int (-1)))

exdiv4c = 
    fmap ("\\x_0 -> GHC.Base.assert (x_0 GHC.Classes.> 0) "++
	  "(GHC.Base.assert ((GHC.Num.+) x_0 (-1) GHC.Classes.> 0) "++
	  "((\\x_1 -> (GHC.Num.+) ((GHC.Num.*) 100 200) "++
	        "(GHC.Real.div x_1 x_0)) "++
	  "((GHC.Num.+) ((GHC.Num.*) 100 200) "++
	  "(GHC.Real.div 5 ((GHC.Num.+) x_0 (-1))))))" ==) $
    runC exdiv4

main = foldr (liftM2 (&&)) (return True) [
       an1r, an1c,
       an2r, an2c,
       an3r, an3c,
       return exdiv1c,
       exdiv2c, exdiv2c',
       exdiv3c,
       exdiv4c
       ]