core/memoize/memoize.factor

   1 ! Copyright (C) 2007, 2010 Slava Pestov, Daniel Ehrenberg.
   2 ! See https://factorcode.org/license.txt for BSD license.
   3 USING: accessors arrays assocs definitions effects hashtables
   4 kernel kernel.private math sequences sequences.private words ;
   5 IN: memoize
   6
   7 <PRIVATE
   8
   9 ! We can't use n*quot, narray and firstn from generalizations because
  10 ! they're macros, and macros use memoize!
  11 : (n*quot) ( n quot -- quotquot )
  12     <repetition> [ ] concat-as ;
  13
  14 : [nsequence] ( length exemplar -- quot )
  15     [ [ [ 1 - ] keep ] dip '[ _ _ _ new-sequence ] ]
  16     [ drop [ [ set-nth-unsafe ] 2keep [ 1 - ] dip ] (n*quot) ] 2bi
  17     [ nip ] 3append ;
  18
  19 : [firstn] ( length -- quot )
  20     [ 0 swap ] swap
  21     [ [ nth-unsafe ] 2keep [ 1 + ] dip ] (n*quot)
  22     [ 2drop ] 3append ;
  23
  24 : packer ( seq -- quot )
  25     length dup 4 <=
  26     [ { [ f ] [ ] [ 2array ] [ 3array ] [ 4array ] } nth ]
  27     [ { } [nsequence] ] if ;
  28
  29 : unpacker ( seq -- quot )
  30     length dup dup 4 <=
  31     [ { [ drop ] [ ] [ first2-unsafe ] [ first3-unsafe ] [ first4-unsafe ] } nth ]
  32     [ [firstn] ] if swap 1 >
  33     [ [ { array } declare ] prepose ] when ;
  34
  35 : pack/unpack ( quot effect -- newquot )
  36     [ in>> packer ] [ out>> unpacker ] bi surround ;
  37
  38 : unpack/pack ( quot effect -- newquot )
  39     [ in>> unpacker ] [ out>> packer ] bi surround ;
  40
  41 : make/n ( table quot effect -- quot )
  42     [ unpack/pack '[ _ _ cache ] ] keep pack/unpack ;
  43
  44 : make/0 ( table quot effect -- quot )
  45     out>> [
  46         packer '[
  47             _ dup first-unsafe [ second-unsafe ] [
  48                 @ @ [
  49                     1 pick set-nth-unsafe t 0 rot set-nth-unsafe
  50                 ] keep
  51             ] if
  52         ]
  53     ] keep unpacker compose ;
  54
  55 : make-memoizer ( table quot effect -- quot )
  56     dup in>> length zero? [ make/0 ] [ make/n ] if ;
  57
  58 : memo-cache ( effect -- cache )
  59     in>> length zero? [ f f 2array ] [ H{ } clone ] if ;
  60
  61 : identity-memo-cache ( effect -- cache )
  62     in>> length zero? [ f f 2array ] [ IH{ } clone ] if ;
  63
  64 PRIVATE>
  65
  66 : (define-memoized) ( word quot effect hashtable -- )
  67     [ [ drop "memo-quot" set-word-prop ] ] dip
  68     '[ 2drop _ "memoize" set-word-prop ]
  69     [ [ [ dup "memoize" word-prop ] 2dip make-memoizer ] keep define-declared ]
  70     3tri ;
  71
  72 : define-memoized ( word quot effect -- )
  73     dup memo-cache (define-memoized) ;
  74
  75 : define-identity-memoized ( word quot effect -- )
  76     dup identity-memo-cache (define-memoized) ;
  77
  78 PREDICATE: memoized < word "memoize" word-prop >boolean ;
  79
  80 M: memoized definer
  81     def>> 3 from-tail swap ?nth hashtable?
  82     \ MEMO: \ IDENTITY-MEMO: ? \ ; ;
  83
  84 M: memoized definition "memo-quot" word-prop ;
  85
  86 M: memoized reset-word
  87     [ call-next-method ]
  88     [ { "memoize" "memo-quot" } remove-word-props ]
  89     bi ;
  90
  91 : memoize-quot ( quot effect -- memo-quot )
  92     dup memo-cache -rot make-memoizer ;
  93
  94 : identity-memoize-quot ( quot effect -- memo-quot )
  95     dup identity-memo-cache -rot make-memoizer ;
  96
  97 : reset-memoized ( word -- )
  98     "memoize" word-prop dup sequence?
  99     [ f swap set-first ] [ clear-assoc ] if ;
 100
 101 : invalidate-memoized ( inputs... word -- )
 102     [ stack-effect in>> packer call ]
 103     [
 104         "memoize" word-prop dup sequence?
 105         [ f swap set-first ] [ delete-at ] if
 106     ]
 107     bi ;
 108
 109 \ invalidate-memoized t "no-compile" set-word-prop