urbit/pkg/king/lib/Noun/Core.hs

{-# OPTIONS_GHC -funbox-strict-fields #-}
{-# LANGUAGE Strict     #-}
{-# LANGUAGE StrictData #-}

module Noun.Core
  ( Noun, nounSize
  , pattern Cell, pattern Atom
  , pattern C, pattern A
  , textToUtf8Atom, utf8AtomToText
  ) where

import ClassyPrelude hiding (hash)

import Noun.Atom

import Control.Lens              (view, from, (&), (^.))
import Data.Bits                 (xor)
import Data.Hashable             (hash)
import GHC.Natural               (Natural)
import GHC.Prim                  (reallyUnsafePtrEquality#)
import Test.QuickCheck.Arbitrary (Arbitrary(arbitrary))
import Test.QuickCheck.Gen       (Gen, getSize, resize, scale)

import qualified Data.Char          as C
import qualified Data.Text.Encoding as T


-- Types -----------------------------------------------------------------------

data Noun
    = NCell Int Word Noun Noun
    | NAtom Int Atom

pattern Cell x y <- NCell _ _ x y where Cell = mkCell
pattern Atom a   <- NAtom _ a     where Atom = mkAtom

{-# COMPLETE Cell, Atom #-}

pattern C x y <- NCell _ _ x y where C = mkCell
pattern A a   <- NAtom _ a     where A = mkAtom

{-# COMPLETE C, A #-}


--------------------------------------------------------------------------------

instance Hashable Noun where
  hash = \case NCell h _ _ _ -> h
               NAtom h _     -> h
  {-# INLINE hash #-}
  hashWithSalt = defaultHashWithSalt
  {-# INLINE hashWithSalt #-}

textToUtf8Atom :: Text -> Noun
textToUtf8Atom = Atom . view (from atomBytes) . encodeUtf8

utf8AtomToText :: Noun -> Either Text Text
utf8AtomToText = \case
    Cell _ _ -> Left "Expected @t, but got ^"
    Atom atm -> T.decodeUtf8' (atm ^. atomBytes) & \case
        Left err -> Left (tshow err)
        Right tx -> pure tx

instance Show Noun where
  show = \case Atom a   -> showAtom a
               Cell x y -> fmtCell (show <$> (x : toTuple y))
    where
      fmtCell :: [String] -> String
      fmtCell xs = "(" <> intercalate ", " xs <> ")"

      toTuple :: Noun -> [Noun]
      toTuple (Cell x xs) = x : toTuple xs
      toTuple atom        = [atom]

      showAtom :: Atom -> String
      showAtom 0 = "0"
      showAtom a | a >= 2^1024 = "\"...\""
      showAtom a =
          let mTerm = do
                t <- utf8AtomToText (Atom a)
                let ok = \x -> (C.isPrint x)
                if (all ok (t :: Text))
                    then pure ("\"" <> unpack t <> "\"")
                    else Left "Don't show as text."

          in case mTerm of
               Left _   -> show a
               Right st -> st

instance Eq Noun where
  (==) x y =
    case reallyUnsafePtrEquality# x y of
      1# -> True
      _  -> case (x, y) of
              (NAtom x1 a1, NAtom x2 a2) ->
                  x1 == x2 && a1 == a2
              (NCell x1 s1 h1 t1, NCell x2 s2 h2 t2) ->
                  s1==s2 && x1==x2 && h1==h2 && t1==t2
              _ ->
                  False
  {-# INLINE (==) #-}

instance Ord Noun where
  compare x y =
    case reallyUnsafePtrEquality# x y of
      1# -> EQ
      _  -> case (x, y) of
              (Atom _,     Cell _ _)   -> LT
              (Cell _ _,   Atom _)     -> GT
              (Atom a1,    Atom a2)    -> compare a1 a2
              (Cell h1 t1, Cell h2 t2) -> compare h1 h2 <> compare t1 t2
  {-# INLINE compare #-}


instance Arbitrary Noun where
  arbitrary = resize 1000 go
    where
      dub x = Cell x x
      go = do
        sz  <- getSize
        (bit, bat :: Bool) <- arbitrary
        case (sz, bit, bat) of
          ( 0, _,     _    ) -> Atom <$> genAtom
          ( _, False, _    ) -> Atom <$> genAtom
          ( _, True,  True ) -> dub <$> arbitrary
          ( _, True,  _    ) -> scale (\x -> x-10) (Cell <$> go <*> go)

genNatural :: Gen Natural
genNatural = fromInteger . abs <$> arbitrary

genAtom :: Gen Atom
genAtom = do
  arbitrary >>= \case
    False -> genNatural
    True  -> (`mod` 16) <$> genNatural

--------------------------------------------------------------------------------

{-# INLINE nounSize #-}
nounSize :: Noun -> Word
nounSize = \case
  NCell _ s _ _ -> s
  NAtom _ _     -> 1

{-# INLINE mkAtom #-}
mkAtom :: Atom -> Noun
mkAtom a = NAtom (hash a) a

{-# INLINE mkCell #-}
mkCell :: Noun -> Noun -> Noun
mkCell h t = NCell has siz h t
  where
    siz = nounSize h + nounSize t
    has = hash h `combine` hash t


-- Stolen from Hashable Library ------------------------------------------------

{-# INLINE combine #-}
combine :: Int -> Int -> Int
combine h1 h2 = (h1 * 16777619) `xor` h2

{-# INLINE defaultHashWithSalt #-}
defaultHashWithSalt :: Hashable a => Int -> a -> Int
defaultHashWithSalt salt x = salt `combine` hash x