From 38098122f31f4f2f73e96ee761a90d17d3fe8dc8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Frederik=20Hangh=C3=B8j=20Iversen?= <fiversen@zendesk.com>
Date: Fri, 11 Oct 2019 16:55:48 +0200
Subject: [PATCH] Allow applications to appear in function OR module contexts

---
 doc/presentation.org      |  2 ++
 ruby/test.rb              |  1 +
 src/Data/Language/Ruby.hs |  8 +++----
 src/Rubyhs/References.hs  | 47 +++++++++++++++++++++++----------------
 4 files changed, 35 insertions(+), 23 deletions(-)

diff --git a/doc/presentation.org b/doc/presentation.org
index 78ce7d9..f2238d1 100644
--- a/doc/presentation.org
+++ b/doc/presentation.org
@@ -51,3 +51,5 @@ So for instance.  The monomorphic type `Parser Function` is a
 procedure that generates a value of type `Function`:
 *** Live coding
 Define an instance for `FromJSON` for `Function`.
+
+Define an instance for `FromJSON` for `Definition`.
diff --git a/ruby/test.rb b/ruby/test.rb
index 6f603dd..be77365 100644
--- a/ruby/test.rb
+++ b/ruby/test.rb
@@ -25,6 +25,7 @@ module M
 end
 
 module K
+  f
 end
 
 f(2)
diff --git a/src/Data/Language/Ruby.hs b/src/Data/Language/Ruby.hs
index ddb6163..b1671a4 100644
--- a/src/Data/Language/Ruby.hs
+++ b/src/Data/Language/Ruby.hs
@@ -61,10 +61,10 @@ instance ToJSON Definition where
       x -> x
 
 instance FromJSON Definition where
-  parseJSON val
-    =   (DefModule <$> parseJSON val)
-    <|> (DefFunction <$> parseJSON val)
-    <|> (DefSend <$> parseJSON val)
+  parseJSON v
+    =   (DefModule <$> parseJSON v)
+    <|> (DefFunction <$> parseJSON v)
+    <|> (DefSend <$> parseJSON v)
 
 newtype Namespace = Namespace [Name]
 
diff --git a/src/Rubyhs/References.hs b/src/Rubyhs/References.hs
index 3267ae7..716d2d6 100644
--- a/src/Rubyhs/References.hs
+++ b/src/Rubyhs/References.hs
@@ -8,6 +8,7 @@ module Rubyhs.References
   , Env(..)
   , Result(..)
   , Namespace(..)
+  , Context(..)
   ) where
 
 import Frelude
@@ -21,8 +22,6 @@ import Data.Aeson ((.=))
 import qualified Data.Aeson as Aeson
 import Data.List
 import Data.HashMap.Strict (HashMap)
-import Data.Text.Prettyprint.Doc.Render.Text (renderStrict)
-import Data.Text.Prettyprint.Doc (layoutCompact)
 import Data.Set (Set)
 import qualified Data.Set as Set
 
@@ -68,21 +67,26 @@ prettyFQN FQN{namespace, name} = case namespace of
 prettyNamespace :: Namespace -> Text
 prettyNamespace = convertString . intercalate "::" . fmap convertString . reverse . toList
 
+data Context = ContextFun FQN | ContextMod Namespace
+
+deriving stock instance Eq Context
+deriving stock instance Ord Context
+
 class Monad m => MyMonad (m :: Type -> Type) where
   declaration :: FQN -> m ()
-  application :: Namespace -> FQN -> m ()
-  getContext :: m Namespace
-  writeContext :: Namespace -> m ()
+  application :: Context -> FQN -> m ()
+  getContext :: m Context
+  writeContext :: Context -> m ()
 
 data Env = Env
   { declarations :: Set FQN
-  , applications :: Map Namespace (Set FQN)
-  , context      :: Namespace
+  , applications :: Map Context (Set FQN)
+  , context      :: Context
   }
 
 data Result = Result
   { declarations :: Set FQN
-  , applications :: Map Namespace (Set FQN)
+  , applications :: Map Context (Set FQN)
   }
 instance ToJSON Result where
   toJSON Result{declarations,applications} = Aeson.object
@@ -91,13 +95,12 @@ instance ToJSON Result where
     ]
     where
       f x = fromList @(HashMap _ _) $ go <$> toList x
-      go (x, y) = (renderStrict $ layoutCompact $ pretty x, y)
+      go (x, y) = (prettyContext x, y)
 
-instance Semigroup Env where
-  Env a0 a1 a2 <> Env b0 b1 b2 = Env (a0 <> b0) (a1 <> b1) (a2 <> b2)
-
-instance Monoid Env where
-  mempty = Env mempty mempty mempty
+prettyContext :: Context -> Text
+prettyContext = \case
+  ContextFun q -> prettyFQN q
+  ContextMod ns -> prettyNamespace ns
   
 instance MyMonad (State Env) where
   declaration q = modify go
@@ -113,7 +116,7 @@ instance MyMonad (State Env) where
     where
     go env = env { context = q }
 
-updateContext :: MyMonad m => (Namespace -> Namespace) -> m ()
+updateContext :: MyMonad m => (Context -> Context) -> m ()
 updateContext f = getContext >>= \c -> writeContext (f c)
 
 locally :: MyMonad m => m a -> m a
@@ -126,7 +129,8 @@ locally act = do
 appendToContext :: MyMonad m => Name -> m ()
 appendToContext n = updateContext go
   where
-  go q = name2ns n <> q
+  go (ContextMod q) = ContextMod $ name2ns n <> q
+  go _ = error "Cannot append module to context in function context"
 
 name2ns :: Name -> Namespace
 name2ns (Name o) = go o
@@ -144,7 +148,8 @@ class References a where
 references :: Block -> Result
 references q = Result{declarations, applications}
   where
-  Env{declarations,applications} = execState (entries @_ @(State Env) q) mempty
+  Env{declarations,applications} = execState (entries @_ @(State Env) q) emptyEnv
+  emptyEnv = Env mempty mempty (ContextMod mempty)
 
 instance References Block where
   entries :: forall m . MyMonad m => Block -> m ()
@@ -163,9 +168,13 @@ instance References Module where
 
 instance References Function where
   entries Function{name, block} = do
-    namespace <- getContext
+    namespace <- getContext >>= \case
+      ContextMod c -> pure c
+      _ -> error "Cannot have a function declaration in a function context"
     declaration $ qual namespace name
-    entries block
+    locally $ do
+      writeContext (ContextFun $ qual namespace name)
+      entries block
 
 qual :: Namespace -> Name -> FQN
 qual namespace (Name o) = case o of