rubyhs/src/Data/Language/Ruby/AST.hs

{-# LANGUAGE DuplicateRecordFields, OverloadedLists #-}
{-# OPTIONS_GHC -Wall #-}
module Data.Language.Ruby.AST
  ( Args(..)
  , Begin(..)
  , Statement(..)
  , Function(..)
  , Module(..)
  , Name(..)
  , Send(..)
  , Namespace(..)
  , Block(..)
  , Casgn(..)
  , RArray(..)
  , RArgs(..)
  , Anything(..)
  , Sym(..)
  ) where

import Data.Aeson (parseJSON, Value(..), withArray)
import Frelude
import qualified Data.Aeson.Types as Aeson
import qualified Data.Vector as Vector
import Data.Coerce
import Data.Word

kebabCase :: String -> String
kebabCase = Aeson.camelTo2 '-'

aesonOptions :: Aeson.Options
aesonOptions = Aeson.defaultOptions
  { Aeson.sumEncoding = Aeson.ObjectWithSingleField
  , Aeson.constructorTagModifier = kebabCase
  }

newtype Begin = Begin [Statement]

deriving stock instance Show Begin
deriving stock instance Ord Begin
deriving stock instance Eq Begin
deriving stock instance Generic Begin
instance ToJSON Begin where
  toEncoding = Aeson.genericToEncoding aesonOptions
deriving newtype instance Semigroup Begin
deriving newtype instance Monoid Begin

instance FromJSON Begin where
  parseJSON = withArray "Begin" $ \as -> case Vector.toList as of
    (String "begin":xs) -> Begin <$> traverse parseJSON xs
    _ -> Begin . pure <$> parseJSON (Array as)

-- Should be 'expression'
data Statement
  = StmtModule Module
  | StmtFunction Function
  | StmtSend Send
  | StmtBlock Block
  | StmtConst Namespace
  | StmtCasgn Casgn
  | StmtArray RArray
  | StmtSym Sym
  -- TODO Get rid of this
  | StmtAnything Anything

deriving stock instance Show Statement
deriving stock instance Ord Statement
deriving stock instance Eq Statement
deriving stock instance Generic Statement
instance ToJSON Statement where
  toEncoding = Aeson.genericToEncoding opts
    where
    opts = aesonOptions { Aeson.constructorTagModifier = go }
    go = \case
      "StmtModule"   -> "module"
      "StmtFunction" -> "function"
      "StmtSend"     -> "send"
      "StmtBlock"    -> "block"
      "StmtConst"    -> "const"
      "StmtCasgn"    -> "casgn"
      "StmtArray"    -> "array"
      "StmtSym"      -> "sym"
      x              -> x

instance FromJSON Statement where
  parseJSON v
    =   (StmtModule   <$> parseJSON v)
    <|> (StmtFunction <$> parseJSON v)
    <|> (StmtSend     <$> parseJSON v)
    <|> (StmtBlock   <$> parseJSON v)
    <|> (StmtConst    <$> parseJSON v)
    <|> (StmtCasgn    <$> parseJSON v)
    <|> (StmtArray    <$> parseJSON v)
    <|> (StmtAnything <$> parseJSON v)

data Casgn = Casgn
  { name      :: Name
  , statement :: Statement
  }
deriving stock instance Show Casgn
deriving stock instance Ord Casgn
deriving stock instance Eq Casgn
deriving stock instance Generic Casgn
instance ToJSON Casgn where
  toEncoding = Aeson.genericToEncoding aesonOptions
instance FromJSON Casgn where
  parseJSON = withArray "Module" $ \case
    [String "casgn", _, name, statement]
      ->  Casgn
      <$> parseJSON name
      <*> parseJSON statement
    _ -> empty


data RArray = RArray
  { statements :: [Statement]
  }
deriving stock instance Show RArray
deriving stock instance Ord RArray
deriving stock instance Eq RArray
deriving stock instance Generic RArray
instance ToJSON RArray where
  toEncoding = Aeson.genericToEncoding aesonOptions
instance FromJSON RArray where
  parseJSON = withArray "RArray" $ \as -> case Vector.toList as of
    String "array":xs
      ->  RArray
      <$> traverse parseJSON xs
    _ -> empty

newtype Anything = Anything Value

deriving stock instance Show Anything
instance Ord Anything where
  compare = coerce compareValue
deriving stock instance Eq Anything
deriving stock instance Generic Anything
instance ToJSON Anything where
  toEncoding = Aeson.genericToEncoding aesonOptions
deriving newtype instance FromJSON Anything

-- f do |x|
--   expr
-- end
data Block = Block
  { send  :: Send
  , args  :: RArgs
  , begin :: Begin
  }

deriving stock instance Show Block
deriving stock instance Ord Block
deriving stock instance Eq Block
deriving stock instance Generic Block
instance ToJSON Block where
  toEncoding = Aeson.genericToEncoding aesonOptions
instance FromJSON Block where
  parseJSON = withArray "Block" $ \as -> case Vector.toList as of
    (String "block":send:args:begin:[])
      ->  Block
      <$> parseJSON send
      <*> parseJSON args
      <*> parseJSON begin
    _ -> empty

-- | It's super confusing that I've already defined a node in my AST
-- called args.  This one correspond to the AST node with the label
-- "args" as reported by `ruby-parse`.
newtype RArgs = RArgs [Arg]

deriving stock instance Show RArgs
deriving stock instance Ord RArgs
deriving stock instance Eq RArgs
deriving stock instance Generic RArgs
instance ToJSON RArgs where
  toEncoding = Aeson.genericToEncoding aesonOptions
instance FromJSON RArgs where
  parseJSON = withArray "RArgs" $ \as -> case Vector.toList as of
    (String "args":xs) -> RArgs <$> traverse parseJSON xs
    _ -> empty

data Arg = Arg Atom | KWArg Atom

deriving stock instance Show Arg
deriving stock instance Ord Arg
deriving stock instance Eq Arg
deriving stock instance Generic Arg
instance ToJSON Arg where
  toEncoding = Aeson.genericToEncoding opts
    where
    opts = aesonOptions { Aeson.constructorTagModifier = go }
    go = \case
      "KWArg" -> "kwarg"
      "Arg"   -> "arg"
      x       -> x

instance FromJSON Arg where
  parseJSON = withArray "Arg" $ \as -> case Vector.toList as of
    [String "arg"   , symbol] -> Arg   <$> parseJSON symbol
    [String "kwarg" , symbol] -> KWArg <$> parseJSON symbol
    _ -> empty

newtype Atom = Atom Text

deriving stock instance Show Atom
deriving stock instance Ord Atom
deriving stock instance Eq Atom
deriving stock instance Generic Atom
instance ToJSON Atom where
  toEncoding = Aeson.genericToEncoding aesonOptions

instance FromJSON Atom where
  parseJSON = \case
    String s -> pure $ Atom s
    _ -> empty

newtype Namespace = Namespace [Name]

deriving newtype instance Semigroup Namespace
deriving newtype instance Monoid Namespace
deriving stock instance Show Namespace
deriving stock instance Ord Namespace
deriving stock instance Eq Namespace
deriving stock instance Generic Namespace
instance ToJSON Namespace where
  toEncoding = Aeson.genericToEncoding aesonOptions

instance FromJSON Namespace where
  parseJSON = \case
    Null -> pure mempty
    Array [String "const", namespace, String name] -> (<> Namespace [Name $ String name]) <$> parseJSON namespace
    _ -> empty

data Send = Send
  { args :: Args
  , namespace :: Namespace
  , name :: Name
  }

deriving stock instance Show Send
deriving stock instance Ord Send
deriving stock instance Eq Send
deriving stock instance Generic Send
instance ToJSON Send where
  toEncoding = Aeson.genericToEncoding aesonOptions

instance FromJSON Send where
  parseJSON = withArray "Send" $ \ as -> case Vector.toList as of
    (String "send" : namespace : name : args)
      ->  Send
      <$> parseJSON (Array $ Vector.fromList args)
      <*> parseJSON namespace
      <*> parseJSON name
    _ -> empty

data Module = Module
  { name  :: Name
  , begin :: Begin
  }

deriving stock instance Show Module
deriving stock instance Ord Module
deriving stock instance Eq Module
deriving stock instance Generic Module
instance ToJSON Module where
  toEncoding = Aeson.genericToEncoding aesonOptions

instance FromJSON Module where
  parseJSON = withArray "Module" $ \case
    [String "module", name, begin]
      ->  Module
      <$> parseJSON name
      <*> parseMaybe begin
    [String "class", name, _, begin]
      ->  Module
      <$> parseJSON name
      <*> parseMaybe begin
    _ -> empty

parseMaybe :: FromJSON m => Monoid m => Value -> Aeson.Parser m
parseMaybe = \case
  Null -> pure mempty
  x -> parseJSON x

data Function = Function
  { name  :: Name
  , args  :: RArgs
  , begin :: Begin
  }

deriving stock instance Show Function
deriving stock instance Ord Function
deriving stock instance Eq Function
deriving stock instance Generic Function
instance ToJSON Function where
  toEncoding = Aeson.genericToEncoding aesonOptions

data Args = Args [Statement]

deriving stock instance Show Args
deriving stock instance Ord Args
deriving stock instance Eq Args
deriving stock instance Generic Args
instance ToJSON Args where
  toEncoding = Aeson.genericToEncoding aesonOptions

instance FromJSON Args where
  parseJSON = withArray "Args" $ \xs -> Args <$> traverse parseJSON (toList xs)

instance FromJSON Function where
  parseJSON = withArray "Function" $ \case
    [String "def", name, args, begin]
      ->  Function
      <$> parseJSON name
      <*> parseJSON args
      <*> parseMaybe begin
    _ -> empty

newtype Name = Name Value

deriving stock instance Show Name
instance Ord Name where
  compare = coerce compareValue
deriving stock instance Eq Name
deriving newtype instance ToJSON Name
deriving newtype instance FromJSON Name

compareValue :: Aeson.Value -> Aeson.Value -> Ordering
compareValue v0 v1 = case (v0, v1) of
  -- This case is buggy:
  (Object o0, Object o1)    -> compare (fst <$> toList o0) (fst <$> toList o1)
  (Array a0, Array a1)      -> foldMap (uncurry compareValue) $ zip (toList a0) (toList a1)
  (String s0, String s1)    -> compare s0 s1
  (Number n0, Number n1)    -> compare n0 n1
  (Bool b0, Bool b1)        -> compare b0 b1
  (Null, Null)              -> EQ
  _                         -> (compare `on` cons) v0 v1
    where
    -- Enumerate constructors.
    cons :: Aeson.Value -> Word8
    cons = \case
      Object{} -> 0
      Array{}  -> 1
      String{} -> 2
      Number{} -> 3
      Bool{}   -> 4
      Null{}   -> 5

newtype Sym = Sym Atom

deriving stock instance Show Sym
deriving newtype instance Ord Sym
deriving stock instance Eq Sym
deriving stock instance Generic Sym
instance ToJSON Sym where
  toEncoding = Aeson.genericToEncoding aesonOptions
instance FromJSON Sym where
  parseJSON = withArray "Sym" $ \case
    [String "sym", atom] -> Sym <$> parseJSON atom
    _ -> empty