 : (collect-loop-info) ( tail? nodes -- )
     [ tail-calls ] keep [ collect-loop-info* ] 2each ;
 
-: remember-loop-info ( #recursive -- )
-    V{ } clone label-stack get length loop-info boa
-    swap label>> loop-infos get set-at ;
+: remember-loop-info ( label -- )
+    label-stack get length swap loop-heights get set-at ;
 
 M: #recursive collect-loop-info*
     nip
     [
-        [ label-stack [ swap label>> suffix ] change ]
-        [ remember-loop-info ]
-        [ t swap child>> (collect-loop-info) ]
-        tri
+        [
+            label>>
+            [ label-stack [ swap suffix ] change ]
+            [ remember-loop-info ]
+            [ t >>loop? drop ]
+            tri
+        ]
+        [ t swap child>> (collect-loop-info) ] bi
     ] with-scope ;
 
+: current-loop-nesting ( label -- labels )
+    label-stack get swap loop-heights get at tail ;
+
+: disqualify-loop ( label -- )
+    work-list get push-front ;
+
 M: #call-recursive collect-loop-info*
-    label>> loop-infos get at
-    [ label-stack get swap height>> tail recursive-call boa ]
-    [ calls>> ]
-    bi push ;
+    label>>
+    swap [ dup disqualify-loop ] unless
+    dup current-loop-nesting [ loop-calls get push-at ] with each ;
 
 M: #if collect-loop-info*
     children>> [ (collect-loop-info) ] with each ;


: collect-loop-info ( node -- )
    { } label-stack set
    H{ } clone loop-infos set
    t swap (collect-loop-info) ;

! Sub-assoc of loop-infos
SYMBOL: potential-loops

: remove-non-tail-calls ( -- )
    loop-infos get
    [ nip calls>> [ tail?>> ] all? ] assoc-filter
    potential-loops set ;

: (remove-non-loop-calls) ( loop-infos -- )
    f over [
        ! If label X is called from within a label Y that is
        ! no longer a potential loop, then X is no longer a
        ! potential loop either.
        over potential-loops get key? [
            potential-loops get '[ , key? ] all?
            [ drop ] [ potential-loops get delete-at t or ] if
        ] [ 2drop ] if
    ] assoc-each
    [ (remove-non-loop-calls) ] [ drop ] if ;

: remove-non-loop-calls ( -- )
    ! Boolean is set to t if something changed.
    !  We recurse until a fixed point is reached.
    loop-infos get [ calls>> [ nesting>> ] map concat ] assoc-map
    (remove-non-loop-calls) ;

: detect-loops ( nodes -- nodes )
    dup
    collect-loop-info
    remove-non-tail-calls
    remove-non-loop-calls
    potential-loops get [ drop t >>loop? drop ] assoc-each ;

! Copyright (C) 2008 Slava Pestov.
! See http://factorcode.org/license.txt for BSD license.
USING: kernel sequences namespaces assocs accessors fry
compiler.tree dequeues search-dequeues ;
IN: compiler.tree.loop.detection

! A loop is a #recursive which only tail calls itself, and those
! calls are nested inside other loops only. We optimistically
! assume all #recursive nodes are loops, disqualifying them as
! we see evidence to the contrary.

: (tail-calls) ( tail? seq -- seq' )
    reverse [ swap [ and ] keep ] map nip reverse ;
