<commit_before>{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ExistentialQuantification #-}
module Effect.Logic
  ( Logic, logic
  , Nondet, nondet
  , module Data.Could
  , module Effect.Weave
  ) where

import Mop
import Effect.Weave
import Data.Could
import Unsafe.Coerce
import Data.IORef
import System.IO.Unsafe
import Debug.Trace

-- nondeterministic choice with pruning

data Logic k
  = FreshScope (Integer -> k)
  | forall a. Choose Integer [a] (a -> k)
  | Cut Integer

data Nondet k = Nondet (IORef Integer) k

{-# INLINE nondet #-}
nondet :: Uses Nondet fs m => Instruction Nondet fs m
nondet = Nondet (unsafePerformIO (newIORef 0)) $ \fs ->
  let Nondet i k = view fs
      next = unsafePerformIO (modifyIORef i succ)
  in next `seq` return fs


{-# INLINABLE freshScope #-}
freshScope :: Has Logic fs m => PlanT fs m Integer
freshScope = symbol (FreshScope id)

-- a nondeterministic producer
-- use: pythag n = logic $ \ch yield cut ->
--        ch [1..n] >>= \z -> ch [1..z] >>= \x -> ch [x..z] >>= \y -> do
--          when (x*x + y*y /= z*z) (yield (x,y,z))
--          cut
{-# INLINE logic #-}
logic :: forall fs m r b.
         (Has Logic fs m,Has Weave fs m)
      => (   (forall a. [a] -> PlanT fs m a)
          -> (b -> PlanT fs m ())
          -> (forall d. PlanT fs m d)
          -> PlanT fs m ()
         )
      -> Producer fs b m ()
logic l =
  producer $ \yield -> do
    scope <- freshScope
    go scope
      $ l (\as -> symbol (Choose scope as id))
          yield
          (symbol (Cut scope))
  where
    go scope p0 = go' p0
      where
        try :: forall a. [a] -> (a -> PlanT fs m ()) -> PlanT fs m () -> PlanT fs m ()
        try [] bp or = or
        try (a:as) bp or = try' (bp a)
          where
            try' p =
              case p of
                Step sym bp' ->
                  case prj sym of
                    Just x ->
                      case x of
                        Choose i as' _ ->
                          if i == scope
                          then try as' (unsafeCoerce bp') (try as bp or)
                          else Step sym (\b -> try' (bp' b))
                        Cut i ->
                          if i == scope
                          then try as bp or
                          else Step sym (\b -> try' (bp' b))
                    Nothing -> Step sym (\b -> try' (bp' b))
                M m -> M (fmap try' m)
                Pure r -> try as bp or
        go' p =
          case p of
            Step sym bp ->
              case prj sym of
                Just x ->
                  case x of
                    Choose i as _ ->
                      if i == scope
                      then try as (unsafeCoerce bp) (Pure ())
                      else Step sym (\b -> go' (bp b))
                    -- ignore cuts if no choices
                    Cut _ -> Step sym (\b -> go' (bp b))
                Nothing -> Step sym (\b -> go' (bp b))
            M m -> M (fmap go' m)
            Pure r -> Pure r

instance Pair Nondet Logic where
  pair p (Nondet i k) (FreshScope ik) =
    let n = (unsafePerformIO $ readIORef i)
    in trace (show n) $ p k (ik n)
  pair _ _ _ = error "Logic primitive escaped its scope:\n\
                     \\tAttempting to reuse control flow\
                     \ primitives outside of their scope\
                     \ is unsupported."
<commit_msg>Remove Data.Could dependency.<commit_after>{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ExistentialQuantification #-}
module Effect.Logic
  ( Logic, logic
  , Nondet, nondet
  , module Effect.Weave
  ) where

import Mop
import Effect.Weave
import Unsafe.Coerce
import Data.IORef
import System.IO.Unsafe
import Debug.Trace

-- nondeterministic choice with pruning

data Logic k
  = FreshScope (Integer -> k)
  | forall a. Choose Integer [a] (a -> k)
  | Cut Integer

data Nondet k = Nondet (IORef Integer) k

{-# INLINE nondet #-}
nondet :: Uses Nondet fs m => Instruction Nondet fs m
nondet = Nondet (unsafePerformIO (newIORef 0)) $ \fs ->
  let Nondet i k = view fs
      next = unsafePerformIO (modifyIORef i succ)
  in next `seq` return fs


{-# INLINABLE freshScope #-}
freshScope :: Has Logic fs m => PlanT fs m Integer
freshScope = symbol (FreshScope id)

-- a nondeterministic producer
-- use: pythag n = logic $ \ch yield cut ->
--        ch [1..n] >>= \z -> ch [1..z] >>= \x -> ch [x..z] >>= \y -> do
--          when (x*x + y*y /= z*z) (yield (x,y,z))
--          cut
{-# INLINE logic #-}
logic :: forall fs m r b.
         (Has Logic fs m,Has Weave fs m)
      => (   (forall a. [a] -> PlanT fs m a)
          -> (b -> PlanT fs m ())
          -> (forall d. PlanT fs m d)
          -> PlanT fs m ()
         )
      -> Producer fs b m ()
logic l =
  producer $ \yield -> do
    scope <- freshScope
    go scope
      $ l (\as -> symbol (Choose scope as id))
          yield
          (symbol (Cut scope))
  where
    go scope p0 = go' p0
      where
        try :: forall a. [a] -> (a -> PlanT fs m ()) -> PlanT fs m () -> PlanT fs m ()
        try [] bp or = or
        try (a:as) bp or = try' (bp a)
          where
            try' p =
              case p of
                Step sym bp' ->
                  case prj sym of
                    Just x ->
                      case x of
                        Choose i as' _ ->
                          if i == scope
                          then try as' (unsafeCoerce bp') (try as bp or)
                          else Step sym (\b -> try' (bp' b))
                        Cut i ->
                          if i == scope
                          then try as bp or
                          else Step sym (\b -> try' (bp' b))
                    Nothing -> Step sym (\b -> try' (bp' b))
                M m -> M (fmap try' m)
                Pure r -> try as bp or
        go' p =
          case p of
            Step sym bp ->
              case prj sym of
                Just x ->
                  case x of
                    Choose i as _ ->
                      if i == scope
                      then try as (unsafeCoerce bp) (Pure ())
                      else Step sym (\b -> go' (bp b))
                    -- ignore cuts if no choices
                    Cut _ -> Step sym (\b -> go' (bp b))
                Nothing -> Step sym (\b -> go' (bp b))
            M m -> M (fmap go' m)
            Pure r -> Pure r

instance Pair Nondet Logic where
  pair p (Nondet i k) (FreshScope ik) =
    let n = (unsafePerformIO $ readIORef i)
    in trace (show n) $ p k (ik n)
  pair _ _ _ = error "Logic primitive escaped its scope:\n\
                     \\tAttempting to reuse control flow\
                     \ primitives outside of their scope\
                     \ is unsupported."

{'fraction_non_alphanumeric': 0.11340053364957012, 'fraction_numerical': 0.0014823599169878447, 'mean_word_length': 1.2311507936507937, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 99, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1958499', 'n_tokens_mistral': 2181, 'n_tokens_neox': 2040, 'n_words': 812}