diff --git a/src/useLockBodyScroll.ts b/src/useLockBodyScroll.ts index b32a92eae1..601332975e 100644 --- a/src/useLockBodyScroll.ts +++ b/src/useLockBodyScroll.ts @@ -48,47 +48,65 @@ export default !doc : function useLockBody(locked: boolean = true, elementRef?: RefObject) { elementRef = elementRef || useRef(doc!.body); - useEffect(() => { - const body = getClosestBody(elementRef!.current); - if (!body) { - return; + const lock = body => { + const bodyInfo = bodies.get(body); + if (!bodyInfo) { + bodies.set(body, { counter: 1, initialOverflow: body.style.overflow }); + if (isIosDevice) { + if (!documentListenerAdded) { + document.addEventListener('touchmove', preventDefault, { passive: false }); + + documentListenerAdded = true; + } + } else { + body.style.overflow = 'hidden'; + } + } else { + bodies.set(body, { counter: bodyInfo.counter + 1, initialOverflow: bodyInfo.initialOverflow }); } + }; + const unlock = body => { const bodyInfo = bodies.get(body); - - if (locked) { - if (!bodyInfo) { - bodies.set(body, { counter: 1, initialOverflow: body.style.overflow }); + if (bodyInfo) { + if (bodyInfo.counter === 1) { + bodies.delete(body); if (isIosDevice) { - if (!documentListenerAdded) { - document.addEventListener('touchmove', preventDefault, { passive: false }); + body.ontouchmove = null; - documentListenerAdded = true; + if (documentListenerAdded) { + document.removeEventListener('touchmove', preventDefault); + documentListenerAdded = false; } } else { - body.style.overflow = 'hidden'; + body.style.overflow = bodyInfo.initialOverflow; } } else { - bodies.set(body, { counter: bodyInfo.counter + 1, initialOverflow: bodyInfo.initialOverflow }); + bodies.set(body, { counter: bodyInfo.counter - 1, initialOverflow: bodyInfo.initialOverflow }); } + } + }; + + useEffect(() => { + const body = getClosestBody(elementRef!.current); + if (!body) { + return; + } + if (locked) { + lock(body); } else { - if (bodyInfo) { - if (bodyInfo.counter === 1) { - bodies.delete(body); - if (isIosDevice) { - body.ontouchmove = null; - - if (documentListenerAdded) { - document.removeEventListener('touchmove', preventDefault); - documentListenerAdded = false; - } - } else { - body.style.overflow = bodyInfo.initialOverflow; - } - } else { - bodies.set(body, { counter: bodyInfo.counter - 1, initialOverflow: bodyInfo.initialOverflow }); - } - } + unlock(body); } }, [locked, elementRef.current]); + + // clean up, on un-mount + useEffect(() => { + const body = getClosestBody(elementRef!.current); + if (!body) { + return; + } + return () => { + unlock(body); + }; + }, []); };