From 777bd910952496b8bff5fd2795badc32d5c0561b Mon Sep 17 00:00:00 2001
From: Yann Herklotz <git@ymhg.org>
Date: Tue, 14 May 2019 18:46:02 +0100
Subject: Add lens to focus on specific module

---
 src/VeriFuzz/Verilog/AST.hs | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

(limited to 'src/VeriFuzz/Verilog/AST.hs')

diff --git a/src/VeriFuzz/Verilog/AST.hs b/src/VeriFuzz/Verilog/AST.hs
index 52155db..7fa2fbd 100644
--- a/src/VeriFuzz/Verilog/AST.hs
+++ b/src/VeriFuzz/Verilog/AST.hs
@@ -17,6 +17,7 @@ Defines the types to build a Verilog AST.
 {-# LANGUAGE FlexibleInstances          #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE MultiParamTypeClasses      #-}
+{-# LANGUAGE RankNTypes                 #-}
 {-# LANGUAGE TemplateHaskell            #-}
 {-# LANGUAGE TypeFamilies               #-}
 
@@ -132,6 +133,7 @@ module VeriFuzz.Verilog.AST
     , modConnName
     , modExpr
     -- * Useful Lenses and Traversals
+    , aModule
     , getModule
     , getSourceId
     , mainModule
@@ -525,6 +527,19 @@ getSourceId :: Traversal' Verilog Text
 getSourceId = getModule . modId . _Wrapped
 {-# INLINE getSourceId #-}
 
+-- | May need to change this to Traversal to be safe. For now it will fail when
+-- the main has not been properly set with.
+aModule :: Text -> Lens' SourceInfo ModDecl
+aModule t = lens get_ set_
+  where
+    set_ (SourceInfo top main) v =
+        SourceInfo top (main & getModule %~ update t v)
+    update top v m@(ModDecl (Identifier i) _ _ _ _) | i == top  = v
+                                                    | otherwise = m
+    get_ (SourceInfo _ main) = head . filter (f t) $ main ^.. getModule
+    f top (ModDecl (Identifier i) _ _ _ _) = i == top
+
+
 -- | May need to change this to Traversal to be safe. For now it will fail when
 -- the main has not been properly set with.
 mainModule :: Lens' SourceInfo ModDecl
-- 
cgit